1#!/usr/bin/env python
2
3import os
4import re
5import sys
6
7from ripe.atlas.tools.commands.base import Command, Factory
8from ripe.atlas.tools.commands.measure import Factory as BaseFactory
9from ripe.atlas.tools.exceptions import RipeAtlasToolsException
10
11
12class RipeAtlas(object):
13
14    def __init__(self):
15        self.command = None
16        self.args = []
17        self.kwargs = {}
18
19    def _generate_usage(self):
20        usage = "Usage: ripe-atlas <command> [arguments]\n\n"
21        usage += "Commands:\n"
22        longest_command = 0
23        classes = []
24        for c in Command.get_available_commands():
25            if c == "shibboleet":
26                continue
27            cmd_class = Command.load_command_class(c)
28            classes.append(cmd_class)
29            cmd_name = cmd_class.get_name()
30            if len(cmd_name) > longest_command:
31                longest_command = len(cmd_name)
32        for cmd_cls in classes:
33            usage += "\t{} {}\n".format(
34                cmd_cls.get_name().ljust(longest_command + 1),
35                cmd_cls.DESCRIPTION,
36            )
37        usage += (
38            "\nFor help on a particular command, try "
39            "ripe-atlas <command> --help"
40        )
41        return usage
42
43    def _set_base_command(self):
44        """
45        Sets the base command covering cases where we call it with
46        shortcut or asking for help.
47        """
48        caller = os.path.basename(sys.argv[0])
49        shortcut = re.match('^a(ping|traceroute|dig|sslcert|ntp|http)$', caller)
50
51        if shortcut:
52            self.command = "measure"
53            sys.argv.insert(1, self._translate_shortcut(shortcut.group(1)))
54            return
55
56        if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
57            self.command = "help"
58            return
59
60        self.command = sys.argv.pop(1)
61
62    @staticmethod
63    def _translate_shortcut(shortcut):
64        if shortcut == "dig":
65            return "dns"
66        return shortcut
67
68    def autocomplete(self):
69        """
70        This function is highly inspired from Django's own autocomplete
71        manage.py. For more documentation check
72        https://github.com/django/django/blob/1.9.4/django/core/management/__init__.py#L198-L270
73        """
74
75        def print_options(options, curr):
76            """
77            Prints matching with current word available autocomplete options
78            in a formatted way to look good on bash.
79            """
80            sys.stdout.write(' '.join(sorted(filter(lambda x: x.startswith(curr), options))))
81
82        # If we are not autocompleting continue as normal
83        if 'RIPE_ATLAS_AUTO_COMPLETE' not in os.environ:
84            return
85
86        cwords = os.environ['COMP_WORDS'].split()[1:]
87        cword = int(os.environ['COMP_CWORD'])
88
89        try:
90            curr = cwords[cword - 1]
91        except IndexError:
92            curr = ''
93
94        commands = list(Command.get_available_commands())
95
96        # base caller ripe-atlas
97        if cword == 1:
98            print_options(commands, curr)
99        # special measure command
100        elif cword == 2 and cwords[0] == "measure":
101            print_options(BaseFactory.TYPES.keys(), curr)
102        # rest of commands
103        elif cwords[0] in commands:
104            cmd = self.fetch_command_class(cwords[0], cwords)
105            cmd.add_arguments()
106            options = [sorted(s_opt.option_strings)[0] for s_opt in cmd.parser._actions if s_opt.option_strings]
107            previous_options = [x for x in cwords[1:cword - 1]]
108            options = [opt for opt in options if opt not in previous_options]
109            print_options(options, curr)
110
111        sys.exit(1)
112
113    def fetch_command_class(self, command, arg_options):
114        """Fetches the class responsible for the given command."""
115
116        cmd_cls = Command.load_command_class(command)
117
118        if cmd_cls is None:
119            # Module containing the command class wasn't found
120            raise RipeAtlasToolsException("No such command")
121
122        #
123        # If the imported module contains a `Factory` class, execute that
124        # to get the `cmd` we're going to use.  Otherwise, we expect there
125        # to be a `Command` class in there.
126        #
127
128        if issubclass(cmd_cls, Factory):
129            cmd = cmd_cls(arg_options).create()
130        else:
131            cmd = cmd_cls(*self.args, **self.kwargs)
132
133        return cmd
134
135    def main(self):
136
137        self._set_base_command()
138
139        self.autocomplete()
140
141        if self.command == "help":
142            raise RipeAtlasToolsException(self._generate_usage())
143
144        cmd = self.fetch_command_class(self.command, sys.argv)
145        cmd.init_args()
146        cmd.run()
147
148
149if __name__ == '__main__':
150    try:
151        sys.exit(RipeAtlas().main())
152    except RipeAtlasToolsException as e:
153        e.write()
154        raise SystemExit()
155