#!/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"

Future improvements:
- Validate SSH keys more strictly:
    - validate options string
    - validate X509 cert strings
- Implement command line options to:
    - filter keys based on options better (beyond regex)
    - filter keys based on comments better (beyond regex)
    - filter X509 keys based on DN/subject
    - support multiple inputs for conditions
    - add an advanced conditional filter language
"""

from ast import literal_eval
import sys
import errno
import argparse
import re
import base64
import hashlib
import copy
import textwrap

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_NOX509 = r'\b(?:(?:spki|pgp|x509|x509v3)-)?(?:(?:ssh|sign)-(?:rsa|dss|ed25519)|ecdsa-[0-9a-z-]+|rsa2048-sha256)(?:-cert-v01@openssh\.com|\@ssh\.com)?\b'
KEYTYPE_REGEX_LAZY_X509 = r'\bx509(?:v3)?-(?:(?:ssh|sign)-(?:rsa|dss|ed25519)|ecdsa-[0-9a-z-]+|rsa2048-sha256)(?:-cert-v01@openssh\.com|\@ssh\.com)?\b'
X509_WORDDN = r'(?:(?i)(?:Distinguished[ _-]?Name|DN|Subject)[=:]?)'  # case insensitive!
KEY_REGEX = r'(.*)\s*(?:(' + KEYTYPE_REGEX_LAZY_NOX509 + r')\s+(' + REGEX_BASE64_START3NULL + r')\s*(.*)|(' + KEYTYPE_REGEX_LAZY_X509 + r')\s+(' + X509_WORDDN + '.*))'

# Group 1: options
# Branch 1:
#  Group 2: keytype (any, including x509)
#  Group 3: key blob (non-x509), always starts with AAAA (3 nulls in base64), no whitespace!
#  Group 4: comment (non-x509)
# Branch 2:
#  Group 5: keytype (x509)
#  Group 6: x509 WORDDN followed by x509-specific blob or DN, including whitespace
#
# If the keytype is x509v3-*, then the data block can actually be a certificate
# XOR a base64 block.
# The cert specifier is "DN:/OU=.../SN=.../C=.." etc. By implication, this
# EXCLUDEs the use of an comments, as you CANNOT detect when the DN ends.


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)
    if m.group(5) is not None:
        key_type = m.group(5)
        blob = m.group(6)
        comment = None
    return (options, key_type, blob, comment)


def fingerprint_key(keyblob, fingerprint_format='SHA256'):
    """Generate SSH key fingerprints, using the requested format."""
    # Don't try to fingerprint x509 blobs
    if keyblob is None or not keyblob.startswith('AAAA'):
        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, conditions, strict=False):
    # pylint: disable=invalid-name,line-too-long,too-many-locals
    """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)

    # Set up our conditions with their defaults
    key_type = conditions.get('key_type', None)
    key_blob = conditions.get('key_blob', None)
    key_fingerprint = conditions.get('key_fingerprint', None)
    key_options_re = conditions.get('key_options_re', None)
    key_comment_re = conditions.get('key_comment_re', None)

    # Try to detect the fingerprint format
    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)

    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 and \
            candidate_key_comment is not None:
        strict_pass = True
        match = match and \
                key_comment_re.search(candidate_key_comment) is not None
    if key_options_re is not None:
        strict_pass = True
        match = match and \
                key_options_re.search(candidate_key_options) is not None
    if strict:
        return match and strict_pass
    return match


PROG_EPILOG = textwrap.dedent("""\
Strict match will require that at least one condition matched.
Conditions marked with X may not work correctly with X509 authorized_keys lines.
""")
PROG_DESC = 'OpenSSH AuthorizedKeysCommand to read from cached keys file'

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        prog='AUTHKEYCMD',
        description=PROG_DESC,
        epilog=PROG_EPILOG,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        add_help=False)
    # Arguments
    group = parser.add_argument_group('Mandatory arguments')
    group.add_argument('username',
                       metavar='USERNAME',
                       nargs='?',
                       type=str,
                       help='Username')
    group.add_argument('--username',
                       metavar='USERNAME',
                       dest='username_opt',
                       type=str,
                       help='Username (alternative form)')
    # Conditions
    group = parser.add_argument_group('Match Conditions (optional)')
    group.add_argument('--key-type',
                       metavar='KEY-TYPE',
                       type=str,
                       help='Key type')
    group.add_argument('--key-fingerprint',
                       '--key-fp',
                       metavar='KEY-FP',
                       type=str,
                       help='Key fingerprint X')
    group.add_argument('--key-blob',
                       metavar='KEY-BLOB',
                       type=str,
                       help='Key blob (Base64 section) X')
    group.add_argument('--key-comment-re',
                       metavar='REGEX',
                       type=str,
                       help='Regex to match on comments X')
    group.add_argument('--key-options-re',
                       metavar='REGEX',
                       type=str,
                       help='Regex to match on options')
    # Setup parameters:
    group = parser.add_argument_group('Misc settings')
    group.add_argument(
        '--cache-file',
        metavar='FILENAME',
        default=DEFAULT_SSHKEY_CACHE,
        type=argparse.FileType('r'),
        help='Cache file [%s]' % (DEFAULT_SSHKEY_CACHE,),
    )
    group.add_argument('--strict',
                       action='store_true',
                       default=False,
                       help='Strict match required')
    group.add_argument('--help', action='help', default=False, help='This help')
    # Fire it all
    args = parser.parse_args()

    # Handle that we support both variants
    lst = [args.username, args.username_opt]
    cnt = lst.count(None)
    if cnt == 2:
        parser.error('Username was not specified')
    elif cnt == 0:
        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]

    # Strict makes no sense without at least one condition being specified
    if args.strict:
        d = copy.copy(vars(args))
        for k in ['cache_file', 'strict', 'username']:
            d.pop(k, None)
        if not any(v is not None for v in list(d.values())):
            parser.error(
                'At least one condition must be specified with --strict')

    if args.key_comment_re is not None:
        args.key_comment_re = re.compile(args.key_comment_re)
    if args.key_options_re is not None:
        args.key_options_re = re.compile(args.key_options_re)

    try:
        key_conditions = {
            'key_options_re': args.key_options_re,
            'key_type': args.key_type,
            'key_blob': args.key_blob,
            'key_fingerprint': args.key_fingerprint,
            'key_comment_re': args.key_comment_re,
        }
        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,
                                    conditions=key_conditions,
                                    strict=args.strict):
                        print(k)
    except IOError as err:
        if err.errno in [errno.EPERM, errno.ENOENT]:
            pass
        else:
            raise err
