#!/usr/bin/env python3
# Server for the RFC 4819 "public key assistant" subsystem
#
# (c) 2010-2015 Mantas Mikulėnas <grawity@gmail.com>
# Released under the MIT Expat License (./LICENSE)

import os, sys
import tempfile
import struct

import authorized_keys

def trace(*args):
    print(*args, file=sys.stderr)

class Keystore(object):
    pass

class DebugYamlKeystore(Keystore):
    def __init__(self, path=None):
        import yaml

        self.path = path or os.path.expanduser("~/ssh-keys.yaml")

    def load(self):
        import yaml

        keys = {}

        try:
            with open(self.path, "r") as fh:
                data = yaml.load(fh)
        except:
            return keys

        for item in data:
            keys[item["algo"], item["pubkey"]] = item["attributes"]

        return keys

    def save(self, keys):
        import yaml

        data = []

        for (kalgo, kblob), attrs in keys.items():
            data.append({"algo": kalgo, "pubkey": kblob, "attributes": attrs})

        with open(self.path, "w") as fh:
            fh.write("# vim: ft=yaml:nowrap:noet\n")
            yaml.dump(data, fh)

    def list(self):
        keys = self.load()

        for (kalgo, kblob), attrs in keys.items():
            yield kalgo, kblob, attrs

    def has_key(self, kalgo, kblob):
        keys = self.load()

        return (kalgo, kblob) in keys

    def add(self, kalgo, kblob, attrs):
        keys = self.load()

        keys[kalgo, kblob] = attrs
        self.save(keys)

    def remove(self, kalgo, kblob):
        keys = self.load()

        if (kalgo, kblob) in keys:
            del keys[kalgo, kblob]
            self.save(keys)
            return True
        else:
            return False

    @classmethod
    def knows_attribute(self, attr):
        return True

class OpenSSHKeystore(Keystore): # {{{
    def __init__(self, path=None):
        self.path = path or os.path.expanduser("~/.ssh/authorized_keys")

    def load(self):
        keys = {}
        extra_attrs = []

        ATTR_PREFIX = "# attribute: "

        for line in open(self.path, "r"):
            line = line.strip()
            if not line:
                pass
            elif line.startswith(ATTR_PREFIX):
                # hack to store attributes not supported by authorized_keys
                name, value = line[len(ATTR_PREFIX):].split("=", 2)
                extra_attrs.append({
                    "name": name,
                    "value": value,
                    "critical": False,
                })
            elif line.startswith("#:"):
                pass
            elif line.startswith("#"):
                extra_attrs.append({
                    "name": "x-comment",
                    "value": line[1:],
                    "critical": False,
                })
            else:
                key = authorized_keys.PublicKey(line)
                attrs = self.convopt_openssh_to_vandyke(key.options)
                if len(key.comment):
                    attrs.append({
                        "name": "comment",
                        "value": key.comment,
                        "critical": False,
                    })
                attrs += extra_attrs
                extra_attrs = []
                keys[key.algo, key.blob] = attrs
        return keys

    def save(self, keys):
        with open(self.path, "w") as fh:
            for (kalgo, kblob), attrs in keys.items():
                self._append_key(fh, kalgo, kblob, attrs)

    def _append_key(self, fh, kalgo, kblob, attrs):
        key = authorized_keys.PublicKey()
        key.algo = kalgo
        key.blob = kblob
        for attr in attrs:
            if attr["name"] == "comment":
                key.comment = attr["value"]
        key.options, unknown_attrs = self.convopt_vandyke_to_openssh(attrs)
        if len(key.comment):
            fh.write("#: %s\n" % key.comment)
        fpr = ":".join("%02x" % ord(c) for c in key.fingerprint())
        print("#: %s" % fpr, file=fh)
        for attr in unknown_attrs:
            if attr["name"] == "x-comment":
                print("#%s" % attr["value"], file=fh)
            else:
                print("# attr: %(name)s=%(value)s" % attr, file=fh)
        print(key, file=fh)
        print("", file=fh)

    def list(self):
        keys = self.load()

        for (kalgo, kblob), attrs in keys.items():
            yield kalgo, kblob, attrs

    def has_key(self, kalgo, kblob):
        keys = self.load()

        return (kalgo, kblob) in keys

    def add(self, kalgo, kblob, attrs):
        with open(self.path, "a") as fh:
            self._append_key(fh, kalgo, kblob, attrs)

        #keys = self.load()
        #keys[kalgo, kblob] = attrs
        #return self.save(keys)

    def remove(self, kalgo, kblob):
        keys = self.load()

        if (kalgo, kblob) in keys:
            del keys[kalgo, kblob]
            self.save(keys)
            return True
        else:
            return False

    attributes = ("agent", "command-override", "comment", "from",
                  "port-forward", "x11", "x-openssh-option")

    @classmethod
    def knows_attribute(self, name):
        return name in self.attributes

    @staticmethod
    def convopt_openssh_to_vandyke(in_opts):
        tmp_attrs = []

        for opt, value in in_opts:
            if opt == "command":
                tmp_attrs.append(("command-override", value))
            elif opt == "from":
                tmp_attrs.append(("from", value))
            elif opt == "no-agent-forwarding":
                tmp_attrs.append(("agent", ""))
            elif opt == "no-port-forwarding":
                tmp_attrs.append(("port-forward", ""))
                tmp_attrs.append(("reverse-forward", ""))
            elif opt == "no-x11-forwarding":
                tmp_attrs.append(("x11", ""))
            else:
                if value is True:
                    attr_value = opt
                else:
                    attr_value = "%s=%s" % (opt, value)
                tmp_attrs.append(("x-openssh-option", attr_value))

        out_attrs = [{"name": attr[0], "value": attr[1], "critical": False}
                     for attr in tmp_attrs]

        return out_attrs

    @staticmethod
    def convopt_vandyke_to_openssh(in_attrs):
        out_opts = authorized_keys.PublicKeyOptions()
        unknown_attrs = []

        for attr in in_attrs:
            # TODO:
            name, value, _ = attr["name"], attr["value"], attr["critical"]
            if name == "agent":
                out_opts.append(("no-agent-forwarding", True))
            elif name == "command-override":
                out_opts.append(("command", value))
            elif name == "comment":
                pass
            elif name == "from":
                out_opts.append(("from", value))
            elif name == "port-forward":
                out_opts.append(("no-port-forwarding", True))
            elif name == "x11":
                out_opts.append(("no-x11-forwarding", True))
            elif name == "x-openssh-option":
                if "=" in value:
                    out_opts.append(value.split("=", 1))
                else:
                    out_opts.append((value, True))
            else:
                unknown_attrs.append(attr)

        return out_opts, unknown_attrs
# }}}

class SshEndOfStream(Exception):
    pass

class SshStream(object):
    def __init__(self, inputfd, outputfd=None):
        self.inputfd = inputfd
        self.outputfd = outputfd or inputfd

    def read(self, *args):
        buf = self.inputfd.read(*args)
        if not buf:
            raise SshEndOfStream
        return buf

    def write(self, *args):
        return self.outputfd.write(*args)

    def writef(self, *args):
        r = self.write(*args)
        if r:
            r = self.outputfd.flush()
        return r

    def read_u32(self):
        buf = self.read(4)
        val, = struct.unpack("!L", buf)
        return val

    def read_bool(self):
        buf = self.read(1)
        val, = struct.unpack("!?", buf)
        return val

    def read_string(self):
        length = self.read_u32()
        buf = self.read(length)
        return buf

    def read_packet(self):
        packet_length = self.read_u32()
        name = self.read_string()
        data_length = packet_length - (4 + len(name))
        return name, data_length

    def write_packet(self, *data):
        fmt = "!L"
        packed = []
        for datum in data:
            if isinstance(datum, int):
                fmt += "L"
                packed += [datum]
            elif isinstance(datum, bool):
                fmt += "?"
                packed += [datum]
            elif isinstance(datum, str):
                buf = datum.encode("utf-8")
                fmt += "L%ds" % len(buf)
                packed += [len(buf), buf]
            elif isinstance(datum, bytes):
                fmt += "L%ds" % len(datum)
                packed += [len(datum), datum]
            else:
                trace("write_packet(): unknown type %r of %r" % (type(datum), datum))
        data_length = struct.calcsize(fmt) - 4
        buf = struct.pack(fmt, data_length, *packed)
        self.writef(buf)

    def write_status(self, code, msg):
        self.write_packet("status", code, msg, "en_US")
# }}}

class PublicKeySubsystem(object):
    SUCCESS                 = 0
    ACCESS_DENIED           = 1
    STORAGE_EXCEEDED        = 2
    VERSION_NOT_SUPPORTED   = 3
    KEY_NOT_FOUND           = 4
    KEY_NOT_SUPPORTED       = 5
    KEY_ALREADY_PRESENT     = 6
    GENERAL_FAILURE         = 7
    REQUEST_NOT_SUPPORTED   = 8
    ATTRIBUTE_NOT_SUPPORTED = 9

    statuses = {
        SUCCESS:                 "Success",
        ACCESS_DENIED:           "Access denied",
        VERSION_NOT_SUPPORTED:   "Protocol version not supported",
        KEY_NOT_FOUND:           "Key not found",
        KEY_NOT_SUPPORTED:       "Key type not supported",
        KEY_ALREADY_PRESENT:     "Key already present",
        GENERAL_FAILURE:         "General failure",
        REQUEST_NOT_SUPPORTED:   "Request not supported",
        ATTRIBUTE_NOT_SUPPORTED: "Attribute not supported",
    }

    fatal_codes = { VERSION_NOT_SUPPORTED }

    def __init__(self, stream):
        self.keystore = None
        self.stream = stream

    def recv_list(self):
        for kalgo, kblob, attrs in self.keystore.list():
            data = [kalgo, kblob, len(attrs)]
            for attr in attrs:
                data += attr["name"], attr["value"]
            self.stream.write_packet("publickey", *data)

        return self.SUCCESS

    def recv_add(self):
        kalgo = self.stream.read_string()
        kblob = self.stream.read_string()
        overwrite = self.stream.read_bool()
        num_attrs = self.stream.read_u32()
        attrs = []

        while num_attrs:
            attrs.append({
                "name": self.stream.read_string(),
                "value": self.stream.read_string(),
                "critical": self.stream.read_bool(),
            })
            num_attrs -= 1

        if self.keystore.has_key(kalgo, kblob) and not overwrite:
            return self.KEY_ALREADY_PRESENT

        for attr in attrs:
            if attr["critical"] and not keystore.knows_attribute(attr["name"]):
                return self.ATTRIBUTE_NOT_SUPPORTED

        self.keystore.add(kalgo, kblob, attrs)

        return self.SUCCESS

    def recv_remove(self):
        kalgo = self.stream.read_string()
        kblob = self.stream.read_string()

        if self.keystore.remove(kalgo, kblob):
            return self.SUCCESS
        else:
            return self.KEY_NOT_FOUND

    def recv_listattributes(self):
        for attr in KnownAttributes:
            self.stream.write_packet("attribute", attr, False)

        return self.SUCCESS

    def handle_message(self, name, data_length):
        if name == b"version":
            ver = self.stream.read_u32()
            if ver == 2:
                self.stream.write_packet("version", ver)
                return None
            else:
                return self.VERSION_NOT_SUPPORTED
        elif name == b"list":
            return self.recv_list()
        elif name == b"add":
            return self.recv_add()
        elif name == b"remove":
            return self.recv_remove()
        elif name == b"listattributes":
            return self.recv_listattributes()
        else:
            trace("received unknown message %r [%r bytes]" % (name, data_length))
            sys.stdin.read(data_length)
            return self.REQUEST_NOT_SUPPORTED

    def loop(self):
        while True:
            name, data_length = self.stream.read_packet()
            code = self.handle_message(name, data_length)
            if code is not None:
                msg = self.statuses.get(code, self.GENERAL_FAILURE)
                self.stream.write_status(code, msg)
                if code in self.fatal_codes:
                    break

    def run(self):
        try:
            self.loop()
        except SshEndOfStream:
            trace("end of stream")
            sys.exit()

def is_interactive():
    try:
        return os.isatty(sys.stdin.fileno())
    except AttributeError:
        return false

if is_interactive():
    trace("This tool is intended to be run as a SSH subsystem, not interactively.")
    sys.exit(2)

os.umask(0o077)
sys.stdin = open("/dev/stdin", "rb")
sys.stdout = open("/dev/stdout", "wb")
sys.stderr = open("/tmp/publickey-%d.trace" % os.getuid(), "a")

stream = SshStream(sys.stdin, sys.stdout)

pks = PublicKeySubsystem(stream)
pks.keystore = OpenSSHKeystore()
#pks.keystore = DebugYamlKeystore()
pks.run()
