Add a real UI to sshwot-export-known-hosts

src/main-export-known-hosts.py

@@ -1,28 +1,82 @@
+import argparse
 import sys
 
 import process_known_hosts
 import write_file
 
 def main():
-	# TODO: Handle errors
-	# TODO: Add a switch for whether you want to include IPs
-	with open(sys.argv[1], 'r') as f:
-		try:
-			known_host_entries = process_known_hosts.process_file(f)
-		except Exception as err:
-			print('Error: %s' % err, file=sys.stderr)
-			sys.exit(1)
+	parser = argparse.ArgumentParser(
+		description = """Export .ssh/known_hosts as a sshwot file.
+		By default the output is written to stdout.""",
+		# We want to provide help on --help, but the default thing
+		# also adds -h, which we don't want
+		add_help = False
+	)
+
+	# --help to get help
+	parser.add_argument('--help',
+		action = 'help',
+		help = 'show this help message and exit'
+	)
+
+	# -i and --include-ips can be used to not ignore the IPs in the file
+	parser.add_argument('-i', '--include-ips',
+		# We store a constant into a property depending on whether
+		# this argument is used or not
+		action = 'store_const',
+		# The property is ignore_ips (to match the argument of
+		# process_known_hosts.process_file()). Store False if this
+		# switch is used and True otherwise
+		dest = 'ignore_ips', const = False, default = True,
+		help = 'includes IPs in addition to domain names'
+	)
+
+	# -o can be used to write the .sshwot thing to a file instead of stdout
+	parser.add_argument('-o',
+		# We store one argument given after this one to the property
+		# outfile, which will be None otherwise
+		action = 'store',
+		dest = 'outfile',
+		# This is what will be displayed in the help after -o
+		metavar = 'outfile',
+		help = 'write the sshwot file to a given file instead of the stdout'
+	)
+
+	# The required argument
+	parser.add_argument('infile',
+		help = 'a file in .ssh/known_hosts format'
+	)
+
+	# This automatically parses the command line args for us. If it
+	# returns, we have correct arguments
+	args = parser.parse_args()
+
+	try:
+		with open(args.infile, 'r') as f:
+			known_host_entries = process_known_hosts.process_file(f, args.ignore_ips)
+	except Exception as err:
+		print('Error: %s' % err, file=sys.stderr)
+		sys.exit(1)
 
 	# Convert to the entry format for .sshwot files
 	entries = []
 	for known_hosts_entry in known_host_entries:
 		entries.append(process_known_hosts.known_hosts_to_entry(known_hosts_entry))
 
-	# Write to stdout by default
-	# TODO: Add a way to change it
-	# We use sys.stdout.buffer instead of just sys.stdout because we
-	# are writing binary data
-	write_file.write(sys.stdout.buffer, entries)
+	try:
+		if args.outfile is None:
+			# Write to stdout by default
+			# We use sys.stdout.buffer instead of just sys.stdout because we
+			# are writing binary data
+			write_file.write(sys.stdout.buffer, entries)
+
+		else:
+			with open(args.outfile, 'wb') as outf:
+				write_file.write(outf, entries)
+
+	except Exception as err:
+		print('Error: %s' % err, file=sys.stderr)
+		sys.exit(1)
 
 if __name__ == '__main__':
 	main()