1# (c) Stefan Countryman, 2019 

2 

3""" 

4Provision or destroy a bunch of digitalocean servers. User will be asked for 

5confirmation before creation or destruction of droplets proceeds. 

6""" 

7 

8import sys 

9import logging 

10from fnmatch import fnmatch 

11import digitalocean 

12from llama.classes import optional_env_var 

13 

14LOGGER = logging.getLogger(__name__) 

15 

16try: 

17 input = raw_input # pylint: disable=redefined-builtin 

18except NameError: 

19 pass # python3 already uses "input" 

20 

21 

22DROPLET_ROW_FMT_DICT = { 

23 "name": "{:<31}", 

24 "ip": "{:<15}", 

25 "status": "{:<7}", 

26 "tags": "{:<24}", 

27} 

28DEFAULT_COLUMNS = ('name', 'ip', 'status', 'tags') 

29SNAPSHOT = 'gwhen.com-post-er13' 

30NO_TOKEN_WARNING = """ 

31You need to specify a digitalocean API token to use this ``llama com do``. You 

32can generate one from your control panel on digitalocean.com. This script 

33expects the DIGITALOCEAN environmental variable to be set to your API token. 

34 

35Feel free to put your API token in a ``.bashrc`` (or similar) startup script, 

36but: 

37 

38*DO NOT UNDER ANY CIRCUMSTANCES VERSION CONTROL THAT API KEY!* 

39 

40It is effectively a username/password combination, and if your code repository 

41containing that credential is ever exposed, someone will be able to use your 

42token to spin up as many servers (e.g. for crypto-mining) as they want *ON YOUR 

43DIME*! This is a **VERY COMMON ATTACK**; people will search for such keys on 

44GitHub and BitBucket and find them in plain sight. 

45""" 

46TOKEN = optional_env_var(['DIGITALOCEAN'], NO_TOKEN_WARNING)[0] 

47 

48 

49def droplet_row_fmt(columns): 

50 """Get a format string for a given list of columns. Those columns will be 

51 inserted into the format string in order.""" 

52 return " ".join(DROPLET_ROW_FMT_DICT[c] for c in columns) 

53 

54 

55def droplet_header(columns): 

56 """Return a header row with column names.""" 

57 return droplet_row_fmt(columns).format(*[c.upper() for c in columns]) 

58 

59 

60def get_ssh_keys(): 

61 """Get a list of all SSH keys.""" 

62 return digitalocean.Manager(token=TOKEN).get_all_sshkeys() 

63 

64 

65def get_tags(): 

66 """Get a list of all available Droplet tags.""" 

67 return digitalocean.Manager(token=TOKEN).get_all_tags() 

68 

69 

70def new_droplet(name, image=SNAPSHOT, ssh_keys=None): 

71 """Return a Droplet object with the given name from the given snapshot. Use 

72 the returned Droplet object to actually create a corresponding droplet on 

73 DigitalOcean. Specify ``ssh_keys`` to add to the droplet (default: no 

74 keys).""" 

75 # actually create the droplet by calling ``create`` on this 

76 return digitalocean.Droplet(token=TOKEN, 

77 name=name, 

78 region='nyc3', 

79 image=image, 

80 size_slug='s-4vcpu-8gb', 

81 backups=False, 

82 ssh_keys=ssh_keys) 

83 

84 

85def create_droplets(names, image=SNAPSHOT, keys=None, tags=()): 

86 """Create droplets with the given names from the given snapshot. Asks for 

87 user confirmation before creating. If ``keys`` is specified as a list, use 

88 ssh keys whose name or ID exactly equal one of the given keys. Specify tags 

89 to apply to droplets by name in ``tags``.""" 

90 ssh_keys = get_ssh_keys() 

91 if keys is not None: 

92 ssh_keys = [k for k in ssh_keys if k.id in keys or k.name in keys] 

93 manager = digitalocean.Manager(token=TOKEN) 

94 snapshots = [m for m in manager.get_all_snapshots() 

95 if m.name == image] 

96 image = snapshots[0].id if snapshots else image 

97 droplets = [new_droplet(n, image, ssh_keys) for n in names] 

98 print_ssh_keys(ssh_keys) 

99 print("Droplets to be created:") 

100 for drop in droplets: 

101 print(drop) 

102 # print("Tags to be applied to these droplets: {}".format(tags)) 

103 sys.stdout.write("CREATE droplets? (DEFAULT: n) [y/N]: ") 

104 choice = input().lower() 

105 if choice != "y": 

106 print("Canceling creation.") 

107 return 

108 print("CREATING DROPLETS!") 

109 for drop in droplets: 

110 drop.create() 

111 if tags: 

112 dtags = [t for t in get_tags() if t.name in tags] 

113 print("DROPLETS INITIALIZED, APPLYING TAGS: ", [t.name for t in dtags]) 

114 for tag in dtags: 

115 tag.add_droplets(droplets) 

116 print("DONE.") 

117 

118 

119def destroy_droplets(names, tags=()): 

120 """Destroy droplets matching the given names. If names is an empty list, 

121 destroy all droplets with at least one of the tags contained in `tags`. 

122 Asks for user confirmation before destroying.""" 

123 droplets = digitalocean.Manager(token=TOKEN).get_all_droplets() 

124 if names: 

125 destroy_list = [d for d in droplets 

126 if any(fnmatch(d.name, g) for g in names)] 

127 else: 

128 destroy_list = [d for d in droplets if set(d.tags).intersection(tags)] 

129 print("Droplets to be destroyed:") 

130 print_droplets(destroy_list, DEFAULT_COLUMNS) 

131 # prompt the user for confirmation 

132 sys.stdout.write("DESTROY droplets? Can't be undone! (DEFAULT: n) [y/N]: ") 

133 choice = input().lower() 

134 if choice == "y": 

135 print("DESTROYING DROPLETS!") 

136 for droplet in destroy_list: 

137 droplet.destroy() 

138 else: 

139 print("Canceling destruction.") 

140 

141 

142def print_ssh_keys(ssh_keys, header_rows=True): 

143 """Print out a list of ``ssh_keys`` in a nice tabular format.""" 

144 key_fmt = "{:<12} {:<50}" 

145 if header_rows: 

146 print("SSH Keys:\n") 

147 print(key_fmt.format("ID", "NAME")) 

148 print("="*len(key_fmt.format("", ""))) 

149 for key in ssh_keys: 

150 print(key_fmt.format(key.id, key.name)) 

151 

152 

153def print_droplets(droplets, columns, header_rows=True): 

154 """Print out a list of droplets in a nice format.""" 

155 if header_rows: 

156 header = droplet_header(columns) 

157 print("{}\n{}".format(header, "="*len(header))) 

158 fmt = droplet_row_fmt(columns) 

159 for drop in droplets: 

160 fmt_args = { 

161 'name': drop.name, 

162 'ip': drop.ip_address, 

163 'status': drop.status, 

164 'tags': str(drop.tags), 

165 } 

166 # can't left-align lists 

167 print(fmt.format(*[fmt_args[c] for c in columns]))