1# (c) Stefan Countryman, 2019 

2 

3import argparse 

4from fnmatch import fnmatch 

5import digitalocean 

6from llama.com.do import ( 

7 __doc__, 

8 SNAPSHOT, 

9 NO_TOKEN_WARNING, 

10 TOKEN, 

11 DROPLET_ROW_FMT_DICT, 

12 DEFAULT_COLUMNS, 

13 print_droplets, 

14 create_droplets, 

15 destroy_droplets, 

16 get_ssh_keys, 

17 print_ssh_keys, 

18 get_tags, 

19) 

20 

21 

22def get_parser(): 

23 """Parse command line arguments and return them as an argparse 

24 namespace.""" 

25 parser = argparse.ArgumentParser(description=__doc__) 

26 arg = parser.add_argument 

27 arg("-c", "--create", nargs="*", default=(), metavar="DROPLET_NAME", 

28 help="""A list of names of snapshots to create. Will raise an error if 

29 any of these snapshots exists.""") 

30 arg("-d", "--destroy", nargs="*", metavar="DROPLET_NAME", help=""" 

31 A list of droplets to destroy. Can also provide a list of UNIX-style 

32 globs, like `llama-parallel*`, to match multiple droplets following a 

33 pattern. If no droplets are specified, instead provide a list of tags 

34 with the `-t` flag; any droplets with any of these tags will be 

35 destroyed. You will be prompted for confirmation before anything is 

36 deleted.""") 

37 arg("-i", "--image", default=SNAPSHOT, metavar="IMAGE_NAME", help=""" 

38 The name of the digitalocean image or user-created snapshot to use for 

39 newly-created droplets. Use `--list-snapshots` to see snapshot options. 

40 DEFAULT: '{}'""".format(SNAPSHOT)) 

41 arg("-k", "--ssh-keys", nargs="*", default=None, metavar="KEY_NAME", 

42 help="""A list of SSH keys to include in newly-created droplets. Use 

43 ``--list-ssh-keys`` to see choices. Specify the SSH keys by ID or by 

44 name. (default: all available keys)""") 

45 arg("-t", "--tags", nargs="*", default=(), metavar="TAG_NAME", help=""" 

46 Tags to use. If destroying droplets and no droplets are specified, 

47 destroy ANY droplets with at least one of the provided tags. If 

48 creating droplets, apply these tags after droplet creation *if* tags 

49 with those names already exist (non-existing tags will be 

50 ignored).""") 

51 arg("-n", "--no-header-rows", action='store_true', help=""" 

52 If this flag is provided, tables will only print their data rows. 

53 Table titles and headers will be omitted. Useful for piping lists of 

54 droplet attributes to another program.""") 

55 arg("-C", "--columns", nargs="*", metavar="COLUMN_NAME", 

56 default=DEFAULT_COLUMNS, choices=tuple(DROPLET_ROW_FMT_DICT), help=""" 

57 Which columns of Droplet metadata to include when printing droplets. 

58 Can be useful if you just want to print out a bunch of IP addresses, 

59 for example. DEFAULT: {}; CHOICES: 

60 {}""".format(', '.join(DEFAULT_COLUMNS), 

61 ', '.join(DROPLET_ROW_FMT_DICT))) 

62 arg("-S", "--list-snapshots", action="store_true", help=""" 

63 If provided, print available user-created snapshots and quit.""") 

64 arg("-K", "--list-ssh-keys", action="store_true", help=""" 

65 If provided, print available SSH keys that can be used on new droplets 

66 and quit.""") 

67 arg("-T", "--list-tags", action="store_true", help=""" 

68 List available droplet tags and quit.""") 

69 arg("-D", "--list-droplets", nargs="*", metavar="DROPLET_NAME", help=""" 

70 If provided, print current user droplets (with their names, IP 

71 addresses, statuses, and tags) and quit. Optionally, provide a list of 

72 droplet names after this flag to only have information about those 

73 droplets printed. Those droplet names can also take shell-style 

74 wildcards, e.g. `llama-parallel-*`, to search for multiple droplets 

75 whose names fit a pattern (matching done with `fnmatch`).""") 

76 return parser 

77 

78 

79def main(): 

80 """Execute the command line interface.""" 

81 parser = get_parser() 

82 args = parser.parse_args() 

83 if TOKEN is None: 

84 parser.error(NO_TOKEN_WARNING) 

85 if args.list_ssh_keys: 

86 print_ssh_keys(get_ssh_keys(), header_rows=(not args.no_header_rows)) 

87 if args.list_tags: 

88 tag_fmt = "{:<14} {:<64}" 

89 if not args.no_header_rows: 

90 print("Tags:\n") 

91 print(tag_fmt.format("NAME", "TOKEN")) 

92 print("="*len(tag_fmt.format("", ""))) 

93 for tag in get_tags(): 

94 print(tag_fmt.format(tag.name, tag.token)) 

95 if args.list_snapshots: 

96 if not args.no_header_rows: 

97 print("Snapshots:\n") 

98 print("NAME"+" "*36 + "\n" + "="*40) 

99 for snapshot in digitalocean.Manager(token=TOKEN).get_all_snapshots(): 

100 print(snapshot.name) 

101 if args.list_droplets is not None: 

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

103 if args.list_droplets: 

104 droplets = [d for d in droplets 

105 if any(fnmatch(d.name, g) for g in args.list_droplets)] 

106 if not args.no_header_rows: 

107 print("Droplets:\n") 

108 print_droplets(droplets, args.columns, 

109 header_rows=(not args.no_header_rows)) 

110 # quit if we printed stuff 

111 if args.list_droplets is not None or args.list_snapshots: 

112 exit() 

113 if args.create: 

114 create_droplets(args.create, image=args.image, keys=args.ssh_keys, 

115 tags=args.tags) 

116 elif args.destroy is not None: 

117 destroy_droplets(args.destroy, args.tags) 

118 

119 

120if __name__ == '__main__': 

121 main()