// NOTICE: This code is for testing & demonstration purposes only.
use bytes;
use memio;
use errors;
use format::ssh;
use io;
use log;
use net;
use net::ssh::agent;
use net::unix;
use os;
use path;
use strings;
use unix::signal;
use unix::signal::{sig};

let running: bool = true;

type identity = struct {
	comment: str,
	privkey: *ssh::key,
	pubkey: []u8,
};

type state = struct {
	identities: []identity,
};

export fn main() void = {
	let state = state { ... };
	const buf = path::init();
	const sockpath = "./socket";

	const listener = unix::listen(sockpath)!;
	defer {
		net::close(listener)!;
		os::remove(sockpath)!;
	};
	os::chmod(sockpath, 0o700)!;
	log::printfln("Listening at {}", sockpath);

	signal::handle(sig::INT, &handle_signal);
	signal::handle(sig::TERM, &handle_signal);

	for (running) {
		const client = match (net::accept(listener)) {
		case errors::interrupted =>
			continue;
		case let err: net::error =>
			log::fatalf("Error: accept: {}", net::strerror(err));
		case let sock: net::socket =>
			yield sock;
		};
		const agent = agent::new(client);
		defer agent::agent_finish(&agent);
		run(&state, &agent);
	};

	for (let i = 0z; i < len(state.identities); i += 1) {
		const ident = state.identities[i];
		ssh::key_free(ident.privkey);
		free(ident.pubkey);
		free(ident.comment);
	};

	log::printfln("Terminated.");
};

fn run(state: *state, agent: *agent::agent) void = {
	for (true) {
		const msg = match (agent::readmsg(agent)) {
		case (io::EOF | agent::error) =>
			break;
		case void =>
			continue;
		case let msg: agent::message =>
			yield msg;
		};
		defer agent::message_finish(&msg);

		const res = match (msg) {
		case agent::request_identities =>
			yield handle_req_ident(state, agent);
		case let msg: agent::add_identity =>
			yield handle_add_ident(state, &msg, agent);
		case let msg: agent::sign_request =>
			yield handle_sign_request(state, &msg, agent);
		case agent::extension =>
			const answer: agent::message = agent::extension_failure;
			agent::writemsg(agent, &answer)!;
		case => abort();
		};
		match (res) {
		case void => yield;
		case agent::error => abort();
		};
	};
};

fn handle_req_ident(
	state: *state,
	agent: *agent::agent,
) (void | agent::error) = {
	let idents: agent::identities_answer = [];
	defer free(idents);

	for (let i = 0z; i < len(state.identities); i += 1) {
		const ident = &state.identities[i];
		append(idents, agent::identity {
			pubkey = ident.pubkey,
			comment = ident.comment,
		});
	};

	const answer: agent::message = idents;
	agent::writemsg(agent, &answer)!;
};

fn handle_add_ident(
	state: *state,
	msg: *agent::add_identity,
	agent: *agent::agent,
) (void | agent::error) = {
	let sink = memio::dynamic();
	ssh::encode_pubkey(&sink, msg.key)!;
	append(state.identities, identity {
		comment = strings::dup(msg.comment),
		privkey = msg.key,
		pubkey = memio::buffer(&sink),
	});
	const answer: agent::message = agent::agent_success;
	agent::writemsg(agent, &answer)?;
	log::printfln("Added key {}", msg.comment);
};

fn handle_sign_request(
	state: *state,
	msg: *agent::sign_request,
	agent: *agent::agent,
) (void | agent::error) = {
	let key: nullable *identity = null;
	for (let i = 0z; i < len(state.identities); i += 1) {
		let ident = &state.identities[i];
		if (bytes::equal(ident.pubkey, msg.key)) {
			key = ident;
			break;
		};
	};
	const key = match (key) {
	case let key: *identity =>
		yield key;
	case null =>
		const answer: agent::message = agent::agent_failure;
		agent::writemsg(agent, &answer)?;
		return;
	};

	let buf = memio::dynamic();
	defer io::close(&buf)!;
	ssh::sign(&buf, key.privkey, msg.data)!;

	const answer: agent::message = agent::sign_response {
		signature = memio::buffer(&buf),
	};
	agent::writemsg(agent, &answer)?;
	log::printfln("Signed challenge with key {}", key.comment);
};

fn handle_signal(sig: sig, info: *signal::siginfo, ucontext: *void) void = {
	running = false;
};