Parse aguments

kishib.py

@@ -1,12 +1,14 @@
 import base64
 import binascii
+import getopt
 import hashlib
+import os.path
 import secrets
 import socket
 import sys
 
 def error(*args):
-	print('%s: Error:' % sys.argv[0], *args, file = sys.stderr)
+	print('%s: Error:' % os.path.basename(sys.argv[0]), *args, file = sys.stderr)
 	sys.exit(1)
 
 def enc_len(length):
@@ -195,30 +197,54 @@ def verify(client_pubkey, server_pubkey):
 			else:
 				error('Could not transfer the keys')
 
+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)
+	if part =='server' or part is None:
+		print('Usage: %s server [-p <port>]' % os.path.basename(sys.argv[0]), file = sys.stderr)
+	sys.exit(1)
+
 def main():
-	# TODO: Actual agument parsing
 	# TODO: Read pubkeys from files
 	# TODO: Write pubkeys to files
-	if sys.argv[1] == 'server':
-		port = int(sys.argv[2])
-		server_pubkey = b'server\n'
+	if len(sys.argv) < 2:
+		usage()
+
+	command = sys.argv[1]
+	opts, fixed = getopt.gnu_getopt(sys.argv[2:], 'p:')
+
+	# TODO: Select an actual port
+	port = 1234
+	for switch, arg in opts:
+		if switch == '-p':
+			try:
+				port = int(arg)
+			except ValueError:
+				error('Port needs to be a number')
+
+	if command == 'server':
+		if len(fixed) != 0:
+			usage('server')
+
+		server_pubkey = b'server'
 
 		client_pubkey = server(server_pubkey, port)
 
 		verify(client_pubkey, server_pubkey)
 
-	elif sys.argv[1] == 'client':
-		host = sys.argv[2]
-		port = int(sys.argv[3])
-		client_pubkey = b'client\n'
+	elif command == 'client':
+		if len(fixed) != 1:
+			usage('client')
+
+		host, = fixed
+		client_pubkey = b'client'
 
 		server_pubkey = client(client_pubkey, host, port)
 
 		verify(client_pubkey, server_pubkey)
 
 	else:
-		print('Usage: %s server PORT')
-		print('Usage: %s client HOST PORT')
+		usage()
 
 if __name__ == '__main__':
 	main()