import argparse
import sys

import entry
import hashing
import open_default_files
import read_file
import write_file

def main():
	parser = argparse.ArgumentParser(
		description = """Search sshwot file(s) for given host and/or
		fingerprint.""",
		# 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'
	)

	# -o can be used to write the results to a file instead of stdout
	parser.add_argument('-o',
		# Given one argument, we open a file of that name and store it
		# to outfile, which will be sys.stdout.buffer otherwise.
		# We use .buffer since we're going to write binary data
		action = 'store',
		dest = 'outfile',
		type = argparse.FileType('wb'),
		default = sys.stdout.buffer,
		# 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'
	)

	# -h/--host for domain, -p/--port for port (default port is 22)
	parser.add_argument('-h', '--host',
		action = 'store',
		dest = 'host',
		help = 'the domain to filter for'
	)
	parser.add_argument('-p', '--port',
		action = 'store',
		dest = 'port',
		# Automatically convert to integer
		type = int,
		help = 'the port associated with the given host'
	)

	# -f/--fingerprint for fingerprint
	parser.add_argument('-f', '--fingerprint',
		action = 'store',
		dest = 'fingerprint',
		help = 'the fingerprint to filter for'
	)

	# Input file(s)
	parser.add_argument('infiles',
		nargs = '*',
		type = argparse.FileType('rb'),
		# The text shown for these in the usage
		metavar = 'sshwot-file',
		help = 'a sshwot file to search'
	)

	# This automatically parses the command line args for us. If it
	# returns, we have correct arguments
	args = parser.parse_args()

	# Ensure we have a host if port was specified
	if args.port is not None and args.host is None:
		print('If you specify a port you need to specify a domain too', file=sys.stderr)
		sys.exit(1)

	# Ensure we're filtering for something
	if args.host is None and args.fingerprint is None:
		print('Specify a host and/or a fingerprint to filter for', file=sys.stderr)
		sys.exit(1)

	# Default to port 22
	port = 22 if args.port is None else args.port

	# Check the validity of the fingerprint and de-base64 it, if we have it
	if args.fingerprint is not None:
		if args.fingerprint[0:7].upper() != 'SHA256:':
			print('We can only handle sha256 fingerprints (starts with SHA256:)')
			sys.exit(1)
		# We encode this, because hashing.base64dec expects bytes
		fingerprint = hashing.base64dec(args.fingerprint[7:].encode())
		# A valid sha256 fingerprint is 32 bytes
		if len(fingerprint) < 32:
			raise Exception('Fingerprint too short')
		elif len(fingerprint) > 32:
			raise Exception('Fingerprint too long')
	
	else:
		fingerprint = None

	# Use the default files if no input files were specified
	if len(args.infiles) == 0:
		infiles = open_default_files.open_all()
	else:
		infiles = args.infiles

	matches = []
	for infile in infiles:
		entries, file_comment = read_file.read(infile)

		# Filter by host if it's present
		if args.host is not None:
			entries = entry.filter_by_host(entries, args.host, port)
		# Filter by fingerprint if it's present
		if fingerprint is not None:
			entries = entry.filter_by_fingerprint(entries, fingerprint)

		matches.extend(entries)

	# Print the matches in sshwot format
	write_file.write(args.outfile, matches)

if __name__ == '__main__':
	try:
		main()
	except Exception as err:
		print('Error: %s' % err, file=sys.stderr)
		sys.exit(1)