#!/usr/bin/python
# vim: ts=4 sts=4 et:
# pylint: disable=invalid-name,line-too-long
"""
OpenSSH AuthorizedKeysCommand: NSSCache input
Copyright 2016 Gentoo Foundation
Written by Robin H. Johnson <robbat2@gentoo.org>
Distributed under the BSD-3 license.

This script returns one or more authorized keys for use by SSH, by extracting
them from a local cache file /etc/sshkey.cache.

Two variants are supported, based on the existing nsscache code:
Format 1:
 username:key1
 username:key2
Format 2:
 username:['key1', 'key2']

Ensure this script is mentioned in the sshd_config like so:
AuthorizedKeysCommand /path/to/nsscache/authorized-keys-command.py

If you have sufficently new OpenSSH, you can also narrow down the search:
AuthorizedKeysCommand /path/to/nsscache/authorized-keys-command.py --username="%u" --key-type="%t" --key-fingerprint="%f" --key-blob="%k"
"""
from __future__ import print_function
from ast import literal_eval
import sys
import errno
import argparse
import re
import base64
import hashlib

DEFAULT_SSHKEY_CACHE = '/etc/sshkey.cache'

REGEX_BASE64 = r'(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?'
# All of the SSH blobs starts with 3 null bytes , which encode to 'AAAA' in base64
REGEX_BASE64_START3NULL = r'AAAA' + REGEX_BASE64
# This regex needs a lot of work
KEYTYPE_REGEX_STRICT = r'\b(?:ssh-(?:rsa|dss|ed25519)|ecdsa-sha2-nistp(?:256|384|521))\b'
# Docs:
# http://www.iana.org/assignments/ssh-parameters/ssh-parameters.xhtml#ssh-parameters-19
# RFC6187, etc
KEYTYPE_REGEX_LAZY = r'\b(?:(?:x509|x509v3|spki|pgp)-)?(?:(?:ssh|sign)-(?:rsa|dss|ed25519)|ecdsa-[0-9a-z-]+)(?:-cert-v01@openssh\.com|\@ssh\.com)?\b'
KEY_REGEX = r'(.*)\s*(' + KEYTYPE_REGEX_LAZY + r')\s+(' + REGEX_BASE64_START3NULL + r')\s*(.*)'
# TODO: If the keytype is x509v3-*, then the data block can actually be a certificate specifier like "DN:/OU=.../SN=.../C=.." etc

def warning(*objs):
    """ Helper function for output to stderr. """
    print("WARNING: ", *objs, file=sys.stderr)

def parse_key(full_key_line):
    """
    Explode an authorized_keys line including options into the various parts.
    """
    #print(KEY_REGEX)
    m = re.match(KEY_REGEX, full_key_line)
    if m is None:
        warning("Failed to match", full_key_line)
        return (None, None, None, None)
    options = m.group(1)
    key_type = m.group(2)
    blob = m.group(3)
    comment = m.group(4)
    return (options, key_type, blob, comment)

def fingerprint_key(keyblob, fingerprint_format='SHA256'):
    """
    Generate SSH key fingerprints, using the requested format.
    """
    if keyblob is None:
        return None
    try:
        binary_blob = base64.b64decode(keyblob)
    except TypeError as e:
        warning(e, keyblob)
        return None
    if fingerprint_format == 'MD5':
        raw = hashlib.md5(binary_blob).digest()
        return 'MD5:'+":".join("{:02x}".format(ord(c)) for c in raw)
    elif fingerprint_format in ['SHA256', 'SHA512', 'SHA1']:
        h = hashlib.new(fingerprint_format)
        h.update(binary_blob)
        raw = h.digest()
        return fingerprint_format+":"+base64.b64encode(raw).rstrip('=')
    return None

def detect_fingerprint_format(fpr):
    """
    Given a fingerprint, try to detect what fingerprint format is used.
    """
    if fpr is None:
        return None
    for prefix in ['SHA256', 'SHA512', 'SHA1', 'MD5']:
        if fpr.startswith(prefix+':'):
            return prefix
    if re.match(r'^(MD5:)?([0-9a-f]{2}:)+[0-9a-f]{2}$', fpr) is not None:
        return 'MD5'
    # Cannot detect the format
    return None

def validate_key(candidate_key, key_type=None, key_fingerprint=None, key_blob=None, key_comment_re=None, key_options_re=None, fingerprint_format=None, strict=False):
    """
    Validate a potential authorized_key line against multiple conditions
    """
    # Explode the key
    (candidate_key_options, candidate_key_type, candidate_key_blob, candidate_key_comment) = parse_key(candidate_key)

    # Try to detect the fingerprint format
    if fingerprint_format is None:
        fingerprint_format = detect_fingerprint_format(key_fingerprint)
    # Force MD5 prefix on old fingerprints
    if fingerprint_format is 'MD5':
        if not key_fingerprint.startswith('MD5:'):
            key_fingerprint = 'MD5:' + key_fingerprint
    # The OpenSSH base64 fingerprints drops the trailing padding, ensure we do
    # the same on provided input
    if fingerprint_format is not 'MD5' and key_fingerprint is not None:
        key_fingerprint = key_fingerprint.rstrip('=')
    # Build the fingerprint for the candidate key
    # (the func does the padding strip as well)
    candidate_key_fingerprint = fingerprint_key(candidate_key_blob, fingerprint_format=fingerprint_format)

    match = True
    strict_pass = False
    if key_type is not None and candidate_key_type is not None:
        strict_pass = True
        match = match and (candidate_key_type == key_type)
    if key_fingerprint is not None and candidate_key_fingerprint is not None:
        strict_pass = True
        match = match and (candidate_key_fingerprint == key_fingerprint)
    if key_blob is not None and candidate_key_blob is not None:
        strict_pass = True
        match = match and (candidate_key_blob == key_blob)
    if key_comment_re is not None:
        # TODO: special validation todo
        match = match and re.match(key_comment_re, candidate_key_comment) is not None
    if key_options_re is not None:
        # TODO: special validation todo
        match = match and re.match(key_options_re, candidate_key_options) is not None
    if strict:
        return match and strict_pass
    return match

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='OpenSSH AuthorizedKeysCommand', epilog="Strict match will require that at least one of keytype/blob/fingerprint matched.")
    parser.add_argument('--cache-file', metavar='FILENAME', type=argparse.FileType('r'), help='Cache file', default=DEFAULT_SSHKEY_CACHE)
    parser.add_argument('--strict', action="store_true", default=False, help='Strict match required')
    parser.add_argument('--key-type', metavar='KEY-TYPE', type=str, help='Key type')
    parser.add_argument('--key-fingerprint', metavar='KEY-FINGERPRINT', type=str, help='Public key fingerprint')
    parser.add_argument('--key-blob', metavar='KEY-BLOB', type=str, help='Public key blob (Base64 section)')
    #parser.add_argument('--key-comment', metavar='REGEX', type=str, help='Regex to match on comments')
    #parser.add_argument('--key-options', metavar='TODO', type=str, help='TODO')
    parser.add_argument('username', metavar='USERNAME', type=str, help='Username', nargs='?')
    parser.add_argument('--username', metavar='USERNAME', type=str, help='Username (alternative form)', dest='username_opt')
    args = parser.parse_args()

    # Handle that we support both variants
    lst = [args.username, args.username_opt]
    if lst.count(None) in [0, 2]:
        parser.error('Username must be specified either as an option XOR argument.')
    else:
        args.username = [x for x in lst if x is not None][0]

    try:
        with args.cache_file as f:
            for line in f:
                (username, key) = line.split(':', 1)
                if username != args.username:
                    continue
                key = key.strip()
                if key.startswith("[") and key.endswith("]"):
                    # Python array, but handle it safely!
                    keys = [i.strip() for i in literal_eval(key)]
                else:
                    # Raw key
                    keys = [key.strip()]
                for k in keys:
                    if validate_key(candidate_key=k, key_type=args.key_type, key_fingerprint=args.key_fingerprint, key_blob=args.key_blob, strict=args.strict):
                        print(k)
    except IOError as err:
        if err.errno in [errno.EPERM, errno.ENOENT]:
            pass
        else:
            raise err