Coverage for src/tagmania/tag_manager.py: 98%

1"""AWS Cluster Tag Management CLI. 

2 

3This module provides a command-line interface for managing tags on AWS resources 

4within a cluster. It supports adding and removing tags from EC2 instances, 

5volumes, snapshots, and Lambda functions belonging to a cluster. 

6 

7The tool processes tag operations in the order specified on the command line, 

8allowing for complex tag management scenarios including tag overrides and 

9conditional operations. 

10 

11Features: 

12 - Add tags to all resources in a cluster 

13 - Remove tags from all resources in a cluster 

14 - Support for multiple resource types (instances, volumes, snapshots, Lambda) 

15 - Batch operations for efficiency 

16 - Tag override capability 

17 

18Usage: 

19 The module is typically invoked via a CLI command for tag management: 

20 

21 ```bash 

22 # Add tags to cluster 

23 tag-manager --tag Environment production --tag Owner backend-team cluster-name 

24 

25 # Remove tags from cluster 

26 tag-manager --untag TempTag --untag BuildNumber cluster-name 

27 

28 # Override existing tag 

29 tag-manager --tag Environment staging cluster-name 

30 ``` 

31 

32Note: 

33 Operations are performed in command-line argument order. If a tag is both 

34 added and removed in the same command, it will be added first then removed. 

35""" 

36 

37import argparse 

38 

39from tagmania.iac_tools import util 

40from tagmania.iac_tools.clusterset import ClusterSet 

41 

42 

43def main(): 

44 """Main entry point for the cluster tag management CLI. 

45 

46 Parses command line arguments and executes tag operations (add/remove) 

47 on all resources within the specified cluster. Supports both EC2 resources 

48 (instances, volumes, snapshots) and Lambda functions. 

49 

50 The function processes tag operations in the order they appear on the 

51 command line, allowing for complex tag management scenarios. 

52 

53 Raises: 

54 SystemExit: On invalid command line arguments. 

55 AWSError: On AWS API failures during tag operations. 

56 """ 

57 parser = argparse.ArgumentParser( 

58 description="AWS cluster tag manager.", 

59 epilog=""" 

60 This tool relies on the "Cluster" and "Owner" tags on instances and 

61 volumes. IAC automation puts this in place. Tags are 

62 dealt with in the order of the parameter arguments: 

63 i.e. if a tag is listed as both a tag and an untag, it will be added 

64 and then removed from the resources. 

65 To override a tag's value, simply add it as a tag and the value will 

66 be overwritten.""", 

67 ) 

68 parser.add_argument( 

69 "-t", 

70 "--tag", 

71 nargs="*", 

72 dest="tag", 

73 help=""" 

74 Add tags to cluster CLUSTER. Tags must take the form of key:value 

75 Can add as many tags as desired (AWS has a 50 tag limit), seperated 

76 by spaces.""", 

77 ) 

78 parser.add_argument( 

79 "-u", 

80 "--untag", 

81 nargs="*", 

82 dest="untag", 

83 help=""" 

84 Remove tags from cluster CLUSTER. Accepts only the key of the tag. 

85 Can remove as many tags as desired, seperated by spaces. If the key 

86 doesn't exist, ignores.""", 

87 ) 

88 parser.add_argument( 

89 "cluster", 

90 help=""" 

91 the name CLUSTER of the cluster in question. This can be found by 

92 looking at any node in the AWS console and looking for the "Cluster" 

93 tag.""", 

94 ) 

95 

96 parser.add_argument("--profile", "-p", help="the AWS profile to use", default=None) 

97 args = parser.parse_args() 

98 

99 cluster = ClusterSet(args.cluster, profile=args.profile) 

100 

101 # Perform tagging operations 

102 if args.tag: 

103 tags = [] 

104 for tag in args.tag: 

105 pair = tag.split(":") 

106 tags.append({"Key": pair[0], "Value": pair[1]}) 

107 

108 # Tag EC2 resources 

109 cluster.tag_instances(tags) 

110 cluster.tag_volumes(tags) 

111 cluster.tag_snapshots(tags) 

112 # cluster.tag_subnet(tags) 

113 

114 # Tag lambda functions 

115 print("Tagging lambda function(s)") 

116 lambda_arn = util.get_lambda_arn("TAGMANIA-AUTO-" + args.cluster) 

117 # Lambda tagging uses an plain dictionary of keys and values 

118 lambda_tags = {} 

119 for tag in tags: 

120 key = tag["Key"] 

121 value = tag["Value"] 

122 lambda_tags[key] = value 

123 util.tag_lambda_functions([lambda_arn], lambda_tags) 

124 

125 # Perform un-tagging operations 

126 if args.untag: 

127 tags = [] 

128 for tag in args.untag: 

129 tags.append({"Key": tag}) 

130 

131 # Un-tag EC2 resources 

132 cluster.untag_instances(tags) 

133 cluster.untag_volumes(tags) 

134 cluster.untag_snapshots(tags) 

135 # cluster.untag_subnet(tags) 

136 

137 # Un-tag lambda functions 

138 print("Un-tagging lambda function(s)") 

139 lambda_arn = util.get_lambda_arn("TAGMANIA-AUTO-" + args.cluster) 

140 # Lambda un-tagging uses a plain list of keys 

141 lambda_keys = [] 

142 for tag in tags: 

143 key = tag["Key"] 

144 lambda_keys.append(key) 

145 util.untag_lambda_functions([lambda_arn], lambda_keys) 

146 

147 

148if __name__ == "__main__": 

149 main()