diff --git a/kishib.py b/kishib.py index 5dcbc7c..a0706d6 100644 --- a/kishib.py +++ b/kishib.py @@ -177,10 +177,13 @@ def parse_hash(auth_hash): return binary # ------------------------------------------------------------------ -# Public key parsing +# File format parsing and serialization # ------------------------------------------------------------------ class PubkeyParseError(Exception): pass def parse_pubkey(pubkey): + # Strip trailing newlines + while pubkey[-1:] in (b'\r', b'\n'): + pubkey = pubkey[:-1] fields = pubkey.split(b' ') # algorithm keymaterial [comment] @@ -195,15 +198,27 @@ def parse_pubkey(pubkey): return algorithm, keymaterial, comment +def serialize_known_hosts(hostname, port, algorithm, keymaterial): + if port is None: + return b'%s %s %s\n' % (hostname, algorithm, keymaterial) + else: + return b'[%s]:%i %s %s\n' % (hostname, port, algorithm, keymaterial) + +def serialize_authorized_keys(algorithm, keymaterial, comment): + if comment is None: + return b'%s %s\n' % (algorithm, keymaterial) + else: + return b'%s %s %s\n' % (algorithm, keymaterial, comment) + # ------------------------------------------------------------------ # UI # ------------------------------------------------------------------ def usage(part = None): if part == 'client' or part is None: - print('Usage: %s client [-p ] [-i ] ' % os.path.basename(sys.argv[0]), file = sys.stderr) + print('Usage: %s client [-p ] [-i ] [-o ] [-P ] ' % os.path.basename(sys.argv[0]), file = sys.stderr) if part =='server' or part is None: - print('Usage: %s server [-p ] [-i ]' % os.path.basename(sys.argv[0]), file = sys.stderr) + print('Usage: %s server [-p ] [-i ] [-o ]' % os.path.basename(sys.argv[0]), file = sys.stderr) sys.exit(1) def verify(client_pubkey, server_pubkey): @@ -236,27 +251,35 @@ def verify(client_pubkey, server_pubkey): error('Could not transfer the keys') def main(): - # TODO: Write pubkeys to files if len(sys.argv) < 2: usage() command = sys.argv[1] - opts, fixed = getopt.gnu_getopt(sys.argv[2:], 'p:i:') + opts, fixed = getopt.gnu_getopt(sys.argv[2:], 'p:i:o:P:') # TODO: Select an actual port port = 1234 pubkey_file = None + output_file = None + ssh_port = None for switch, arg in opts: if switch == '-p': try: port = int(arg) except ValueError: error('Port needs to be a number') - if switch == '-i': + elif switch == '-i': pubkey_file = arg + elif switch == '-o': + output_file = arg + elif switch == '-P': + try: + ssh_port = int(arg) + except ValueError: + error('Port needs to be a number') if command == 'server': - if len(fixed) != 0: + if len(fixed) != 0 or ssh_port is not None: usage('server') if pubkey_file is None: @@ -281,13 +304,26 @@ def main(): with open(pubkey_file, 'rb') as f: server_pubkey = f.read() except IOError as err: - error('Could not open server public key file: %s' % err) + error('Could not read server public key: %s' % err) client_pubkey = server(server_pubkey, port) verify(client_pubkey, server_pubkey) - print(parse_pubkey(client_pubkey)) + algorithm, keymaterial, comment = parse_pubkey(client_pubkey) + authorized_keys_entry = serialize_authorized_keys(algorithm, keymaterial, comment) + + if output_file is None: + # Try ~/.ssh/authorized_keys + if 'HOME' not in os.environ: + error('Cannot locate homedir, $HOME is not set') + output_file = os.environ['HOME'] + '/.ssh/authorized_keys' + + try: + with open(output_file, 'ab') as f: + f.write(authorized_keys_entry) + except IOError as err: + error('Could not write authorized_keys entry: %s' % err) elif command == 'client': if len(fixed) != 1: @@ -318,7 +354,7 @@ def main(): with open(pubkey_file, 'rb') as f: client_pubkey = f.read() except IOError as err: - error('Could not open client public key file: %s' % err) + error('Could not read client public key: %s' % err) host, = fixed # Support internationalized domain names @@ -328,7 +364,20 @@ def main(): verify(client_pubkey, server_pubkey) - print(parse_pubkey(server_pubkey)) + algorithm, keymaterial, comment = parse_pubkey(server_pubkey) + known_hosts_entry = serialize_known_hosts(host.encode(), ssh_port, algorithm, keymaterial) + + if output_file is None: + # Try ~/.ssh/known_hosts + if 'HOME' not in os.environ: + error('Cannot locate homedir, $HOME is not set') + output_file = os.environ['HOME'] + '/.ssh/known_hosts' + + try: + with open(output_file, 'ab') as f: + f.write(known_hosts_entry) + except IOError as err: + error('Could not write known_hosts entry: %s' % err) else: usage()