Read pubkeys from files

kishib.py

@@ -2,7 +2,7 @@ import base64
 import binascii
 import getopt
 import hashlib
-import os.path
+import os
 import secrets
 import socket
 import sys
@@ -176,15 +176,34 @@ def parse_hash(auth_hash):
 
 	return binary
 
+# ------------------------------------------------------------------
+# Public key parsing
+# ------------------------------------------------------------------
+class PubkeyParseError(Exception): pass
+def parse_pubkey(pubkey):
+	fields = pubkey.split(b' ')
+
+	# algorithm keymaterial [comment]
+	if len(fields) < 2:
+		raise PubkeyParseError
+	algorithm, keymaterial, *comment = fields
+
+	if len(comment) == 0:
+		comment = None
+	else:
+		comment = b''.join(comment)
+
+	return algorithm, keymaterial, comment
+
 # ------------------------------------------------------------------
 # UI
 # ------------------------------------------------------------------
 
 def usage(part = None):
 	if part == 'client' or part is None:
-		print('Usage: %s client [-p <port>] <host>' % os.path.basename(sys.argv[0]), file = sys.stderr)
+		print('Usage: %s client [-p <port>] [-i <pubkey>] <host>' % os.path.basename(sys.argv[0]), file = sys.stderr)
 	if part =='server' or part is None:
-		print('Usage: %s server [-p <port>]' % os.path.basename(sys.argv[0]), file = sys.stderr)
+		print('Usage: %s server [-p <port>] [-i <pubkey>]' % os.path.basename(sys.argv[0]), file = sys.stderr)
 	sys.exit(1)
 
 def verify(client_pubkey, server_pubkey):
@@ -217,47 +236,100 @@ def verify(client_pubkey, server_pubkey):
 				error('Could not transfer the keys')
 
 def main():
-	# TODO: Read pubkeys from files
 	# TODO: Write pubkeys to files
 	if len(sys.argv) < 2:
 		usage()
 
 	command = sys.argv[1]
-	opts, fixed = getopt.gnu_getopt(sys.argv[2:], 'p:')
+	opts, fixed = getopt.gnu_getopt(sys.argv[2:], 'p:i:')
 
 	# TODO: Select an actual port
 	port = 1234
+	pubkey_file = 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':
+			pubkey_file = arg
 
 	if command == 'server':
 		if len(fixed) != 0:
 			usage('server')
 
-		server_pubkey = b'server'
+		if pubkey_file is None:
+			# Try the default ssh host key location
+			server_pubkey = None
+
+			for algorithm in ['ed25519', 'ecdsa', 'rsa']:
+				pubkey_file = '/etc/ssh/ssh_host_' + algorithm + '_key.pub'
+
+				try:
+					with open(pubkey_file, 'rb') as f:
+						server_pubkey = f.read()
+						break
+				except IOError:
+					continue
+
+			if server_pubkey is None:
+				error('Could not find server public key (tried /etc/ssh/ssh_host_{ed25519,ecdsa,rsa}_key.pub)')
+
+		else:
+			try:
+				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)
 
 		client_pubkey = server(server_pubkey, port)
 
 		verify(client_pubkey, server_pubkey)
 
+		print(parse_pubkey(client_pubkey))
+
 	elif command == 'client':
 		if len(fixed) != 1:
 			usage('client')
 
+		if pubkey_file is None:
+			# Try the default ssh client key location
+			client_pubkey = None
+
+			if 'HOME' not in os.environ:
+				error('Cannot locate homedir, $HOME is not set')
+
+			for algorithm in ['ed25519', 'ecdsa', 'rsa']:
+				pubkey_file = os.environ['HOME'] +'/.ssh/id_' + algorithm + '.pub'
+
+				try:
+					with open(pubkey_file, 'rb') as f:
+						client_pubkey = f.read()
+						break
+				except IOError:
+					continue
+
+			if client_pubkey is None:
+				error('Could not find client public key (tried ~/.ssh/id_{ed25519,ecdsa,rsa}.pub)')
+
+		else:
+			try:
+				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)
+
 		host, = fixed
 		# Support internationalized domain names
 		host = host.encode('idna').decode()
 
-		client_pubkey = b'client'
-
 		server_pubkey = client(client_pubkey, host, port)
 
 		verify(client_pubkey, server_pubkey)
 
+		print(parse_pubkey(server_pubkey))
+
 	else:
 		usage()