#!/usr/bin/python3
"""Glue to allow a sshd to authenticate users based on keys held in Caelum.

You need two things on each SSH server; this script can be both:
  * A server daemon, which maintains an in-memory set of SSH keys,
    synchronised from a central database (Caelum).
      /path/to/this/script --server wss://sshkeys.example.com/ws
  * A client, which sshd calls to get the keys for a user using
      AuthorizedKeysCommand /path/to/this/script --client "%u" "%f" "%C"
    (Make sure FingerprintHash is sha256, which is generally the default.)
"""

import asyncio
import logging
import json
import os
from collections import defaultdict, UserDict
from functools import reduce

import websockets
import backoff


logger = logging.getLogger("caelum_ssh_authenticator")


async def client(sockpath, owner, fingerprint="", conn_info=""):
	reader, writer = await asyncio.open_unix_connection(sockpath)
	assert "\n" not in owner
	assert "\0" not in owner
	assert "\n" not in fingerprint
	assert "\0" not in fingerprint
	assert "\n" not in conn_info
	assert "\0" not in conn_info
	writer.write(owner.encode()+b"\n")
	writer.write(fingerprint.encode()+b"\n")
	writer.write(conn_info.encode()+b"\n")
	writer.write_eof()
	r = await reader.read()
	print(r.decode(), end="")
	writer.close()
	await writer.wait_closed()


class KeyStore(UserDict):
	def __init__(self):
		# deliberately not calling superconstructor -- that gives us a dict, and we want to shoehorn in a defaultdict
		self.data = defaultdict(lambda: defaultdict(dict))
		self.uplink = None
		self.uplink_ready = asyncio.Future()
		self.log_connection_error = True
		self.newdata = None

	def prepare_replacement(self):
		self.newdata = defaultdict(lambda: defaultdict(dict))
		return self.newdata

	def perform_replacement(self):
		self.data = self.newdata

	def set_uplink_websocket(self, websocket):
		self.uplink = websocket

	async def uplink_send(self, data):
		if self.uplink is not None and self.uplink.close_code is None:
			await self.uplink.send(data)
		else:
			logger.warning("Discarding message to disconnected authentication server: %s", data)


@backoff.on_predicate(backoff.constant, lambda x: x is None, interval=10)
async def uplink(url, keystore):
	try:
		async with websockets.connect(url) as websocket:
			logger.info("Connected to authentication service on %s", url)
			keystore.log_connection_error = True
			keystore.set_uplink_websocket(websocket)
			# Initially write to a new inactive keystore dict, so that after reconnection the old one
			# remains in use for authentication until the new one is fully populated
			writing_keystore = keystore.prepare_replacement()
			while True:
				packet = await websocket.recv()
				logger.debug(packet)
				if packet == "":
					# End of initial burst
					keycount = reduce(lambda cc, dd: cc+reduce(lambda c, d: c+len(d), dd.values(), 0), writing_keystore.values(), 0)
					logger.info("Initial key download: got %d user%s with %d key%s", len(writing_keystore),
						"" if len(writing_keystore)==1 else "s", keycount, "" if keycount==1 else "s")
					keystore.perform_replacement()
					if keystore.uplink_ready is not None and not keystore.uplink_ready.done():
						keystore.uplink_ready.set_result(True)
					continue
				packet = json.loads(packet)
				if "u" not in packet or "i" not in packet or "k" not in packet or "f" not in packet:
					logger.error("Bad packet from authentication service: %s", packet)
					continue
				if packet["k"] is None:
					# Key was deleted.  We know its owner (u) and key ID (i) only.
					# We optimise for fast lookup during auth by indexing on key fingerprint, at the
					# expense of having to scan all of the user's keys when deleting (infrequent)
					dels_f = []
					for cand_f, cand_keylist in writing_keystore[packet["u"]].items():
						for cand_i in cand_keylist.keys():
							if cand_i == packet["i"]:
								dels_f += [cand_f]
					if not dels_f:
						logger.warning("Attempted to delete nonexistent key %s/%s", packet["u"], packet["i"])
					else:
						for del_f in dels_f:
							if packet["t"]: # after initial burst
								logger.info("%s: deleted key %s/%s/%s", packet["t"],
									packet["u"], packet["i"], del_f)
							del writing_keystore[packet["u"]][del_f][packet["i"]]
							if not writing_keystore[packet["u"]][del_f]:
								del writing_keystore[packet["u"]][del_f]
				else:
					if packet["t"]: # after initial burst
						logger.info("%s: got key %s/%s/%s", packet["t"],
							packet["u"], packet["i"], packet["f"])
					writing_keystore[packet["u"]][packet["f"]][packet["i"]] = packet["k"]
		logger.warning("Lost connection to authentication service")
	except (OSError, ConnectionRefusedError) as e:
		if keystore.log_connection_error:
			logger.error("Connection to authentication service failed: %s", e)
			keystore.log_connection_error = False
	except websockets.exceptions.ConnectionClosedError as e:
		logger.error("Lost connection to authentication service: %s", e)
	except websockets.exceptions.InvalidURI as e:
		if keystore.uplink_ready is not None and not keystore.uplink_ready.done():
			keystore.uplink_ready.set_result(False)
		raise e
	except Exception:
		logger.exception("Exception in uplink")


async def handle_client(reader, writer, keystore):
	user = await reader.readline()
	user = user.decode().strip()
	fingerprint = await reader.readline()
	fingerprint = fingerprint.decode().strip()
	conn_info = await reader.read()
	conn_info = conn_info.decode().strip()
	if conn_info:
		_server_ip, _server_port, client_ip, _client_port = conn_info.split()
	else:
		# No support for %C before OpenSSH 9.4
		client_ip = None
	if user in keystore: # don't let defaultdict create dicts for invalid users/keys
		if not fingerprint: # return all keys, for debugging
			logger.info("Listed all keys for %s", user)
			for fingerprint, keys in keystore[user].items():
				for key_id, key in keys.items():
					writer.write(key.encode()+b"\n")
		elif fingerprint in keystore[user]:
			for key_id, key in keystore[user][fingerprint].items():
				logger.info("Authentication request from %s for %s: returned key %s: %s", client_ip, user, key_id, key)
				writer.write(key.encode()+b"\n")
				await keystore.uplink_send(json.dumps({
						"u": user,
						"i": key_id,
						"f": fingerprint,
						"c": client_ip,
					}))
		else:
			logger.warning("Authentication request from %s for %s: no key matching fingerprint %s found", client_ip, user, fingerprint)
	else:
		logger.warning("Authentication request from %s for unknown user %s with fingerprint %s", client_ip, user, fingerprint)
	writer.write_eof()
	writer.close()
	await writer.wait_closed()


async def server(sockpath, url):
	keystore = KeyStore()
	uplink_task = asyncio.create_task(uplink(url, keystore))
	if not await keystore.uplink_ready:
		# uplink_task caught and re-raised an exception
		await uplink_task
		# we probably never get here as 'await uplink_task' would raise its exception, but just in case:
		raise uplink_task.exception()
	s = await asyncio.start_unix_server(lambda r, w: handle_client(r, w, keystore), sockpath)
	async with s:
		logger.info("Server listening on %s", sockpath)
		await s.serve_forever()


if __name__ == "__main__":
	import argparse

	logging.basicConfig(format="%(message)s")
	logger.setLevel(logging.INFO)
	logging.getLogger().setLevel(logging.INFO) # root logger
	logging.getLogger("backoff").setLevel(logging.WARNING)

	parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
	parser.add_argument("--server", "-s", metavar="URL", help="Serve clients on a Unix socket, authenticating data against the Caelum authentication server at ws(s):// URL")
	parser.add_argument("--client", "-c", nargs="+", metavar=('PARAM'), help="Use from AuthorizedKeysCommand ... \"%%u\" \"%%f\" \"%%C\" in sshd_config, or just pass a username to list all their keys")
	parser.add_argument("--socket", "-S", metavar="PATH", default="/run/caelum-ssh-auth/caelum_ssh_authenticator.sock", help="Path to server Unix socket")
	parser.add_argument("--umask", "-U", type=int, default=0o007, help="Umask to restrict access to the server socket")
	args = parser.parse_args()

	if not args.server and not args.client:
		parser.print_help()
		parser.exit(1)

	if args.umask is not None:
		logger.debug("Umask: %s", oct(args.umask))
		os.umask(args.umask)

	try:
		if args.server:
			asyncio.run(server(args.socket, args.server))
		elif args.client:
			asyncio.run(client(args.socket, *args.client))
	except Exception as e:
		logger.critical(e)


# vim: tabstop=4 shiftwidth=4 noexpandtab
