Algorithm lookup (#53)

* Adding ssh-audit.py to algorithm_lookup_branch

* Removed the use of an error handler from algorithm_lookup and implemented suggestions made by jugmac00 and jtesta

README.md

@@ -36,6 +36,7 @@ usage: ssh-audit.py [options] <host>
    -j,  --json             JSON output
    -l,  --level=<level>    minimum output level (info|warn|fail)
    -L,  --list-policies    list all the official, built-in policies
+        --lookup=<alg>     performs an algorithm lookup (accepts a comma separated list)
    -M,  --make-policy=<policy.txt>  creates a policy based on the target server
                                     (i.e.: the target server has the ideal
                                     configuration that other servers should

ssh-audit.py

@@ -79,6 +79,7 @@ def usage(err: Optional[str] = None) -> None:
     uout.info('   -j,  --json             JSON output')
     uout.info('   -l,  --level=<level>    minimum output level (info|warn|fail)')
     uout.info('   -L,  --list-policies    list all the official, built-in policies')
+    uout.info('        --lookup=<alg>     performs an algorithm lookup (accepts a comma separated list)')
     uout.info('   -M,  --make-policy=<policy.txt>  creates a policy based on the target server\n                                    (i.e.: the target server has the ideal\n                                    configuration that other servers should\n                                    adhere to)')
     uout.info('   -n,  --no-colors        disable colors')
     uout.info('   -p,  --port=<port>      port to connect')
@@ -477,6 +478,7 @@ class AuditConf:
         self.target_file = None  # type: Optional[str]
         self.target_list = []  # type: List[str]
         self.list_policies = False
+        self.lookup = ''
 
     def __setattr__(self, name: str, value: Union[str, int, float, bool, Sequence[int]]) -> None:
         valid = False
@@ -518,7 +520,7 @@ class AuditConf:
             if value == -1.0:
                 raise ValueError('invalid timeout: {}'.format(value))
             valid = True
-        elif name in ['policy_file', 'policy', 'target_file', 'target_list']:
+        elif name in ['policy_file', 'policy', 'target_file', 'target_list', 'lookup']:
             valid = True
 
         if valid:
@@ -530,7 +532,7 @@ class AuditConf:
         aconf = cls()
         try:
             sopts = 'h1246M:p:P:jbcnvl:t:T:L'
-            lopts = ['help', 'ssh1', 'ssh2', 'ipv4', 'ipv6', 'make-policy=', 'port=', 'policy=', 'json', 'batch', 'client-audit', 'no-colors', 'verbose', 'level=', 'timeout=', 'targets=', 'list-policies']
+            lopts = ['help', 'ssh1', 'ssh2', 'ipv4', 'ipv6', 'make-policy=', 'port=', 'policy=', 'json', 'batch', 'client-audit', 'no-colors', 'verbose', 'level=', 'timeout=', 'targets=', 'list-policies', 'lookup=']
             opts, args = getopt.gnu_getopt(args, sopts, lopts)
         except getopt.GetoptError as err:
             usage_cb(str(err))
@@ -578,10 +580,15 @@ class AuditConf:
                 aconf.target_file = a
             elif o in ('-L', '--list-policies'):
                 aconf.list_policies = True
+            elif o in ('--lookup'):
+                aconf.lookup = a
 
-        if len(args) == 0 and aconf.client_audit is False and aconf.target_file is None and aconf.list_policies is False:
+        if len(args) == 0 and aconf.client_audit is False and aconf.target_file is None and aconf.list_policies is False and aconf.lookup == '':
             usage_cb()
 
+        if aconf.lookup != '':
+            return aconf
+
         if aconf.list_policies:
             list_policies()
             sys.exit(PROGRAM_RETVAL_GOOD)
@@ -3714,6 +3721,59 @@ def audit(aconf: AuditConf, sshv: Optional[int] = None, print_target: bool = Fal
     return program_retval
 
 
+def algorithm_lookup(alg_names: str) -> int:
+    alg_types = {
+        'kex': 'key exchange algorithms',
+        'key': 'host-key algorithms',
+        'mac': 'message authentication code algorithms',
+        'enc': 'encryption algorithms (ciphers)'
+    }
+
+    algorithm_names = alg_names.split(",")
+    adb = SSH2.KexDB.ALGORITHMS
+
+    # Use nested dictionary comprehension to iterate an outer dictionary where
+    # each key is an alg type that consists of a value (which is itself a
+    # dictionary) of alg names. Filter the alg names against the user supplied
+    # list of names.
+    algorithms_dict = {
+        outer_k: {
+            inner_k
+            for (inner_k, inner_v) in outer_v.items()
+            if inner_k in algorithm_names
+        }
+        for (outer_k, outer_v) in adb.items()
+    }
+
+    unknown_algorithms = []  # type: List[str]
+    padding = len(max(algorithm_names, key=len))
+
+    for alg_type in alg_types:
+        if len(algorithms_dict[alg_type]) > 0:
+            title = str(alg_types.get(alg_type))
+            retval = output_algorithms(title, adb, alg_type, algorithms_dict[alg_type], unknown_algorithms, False, PROGRAM_RETVAL_GOOD, padding)
+
+    algorithms_dict_flattened = [
+        alg_name
+        for val in algorithms_dict.values()
+        for alg_name in val
+    ]
+
+    algorithms_not_found = [
+        alg_name
+        for alg_name in algorithm_names
+        if alg_name not in algorithms_dict_flattened
+    ]
+
+    if len(algorithms_not_found) > 0:
+        retval = PROGRAM_RETVAL_FAILURE
+        out.head('# unknown algorithms')
+        for algorithm_not_found in algorithms_not_found:
+            out.fail(algorithm_not_found)
+
+    return retval
+
+
 utils = Utils()
 out = Output()
 
@@ -3721,6 +3781,10 @@ out = Output()
 def main() -> int:
     aconf = AuditConf.from_cmdline(sys.argv[1:], usage)
 
+    if aconf.lookup != '':
+        retval = algorithm_lookup(aconf.lookup)
+        sys.exit(retval)
+
     # If multiple targets were specified...
     if len(aconf.target_list) > 0:
         ret = PROGRAM_RETVAL_GOOD