#!/usr/bin/env python3
#
## License
#
# Copyright (c) 2020 Jesse Weaver.
#
# This file is part of secretgarden.
# 
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

## Introduction
# `secretgarden` is a simple way to generate and securely store secrets for sysadmins that manage a
# small set of systems by themselves. It can be used as a standalone CLI tool, or as a lookup plugin
# within Ansible.
#
# It is strongly inspired by the interaction model of `BOSH` and `CredHub`, where an automated
# deployment tool can ask for a secret that is:
#   - Securely stored
#   - Automatically generated if it does not exist
#   - Re-generated if the generation options have changed
#   - Easy to generate based on other certificates (CAs with child certificates, for instance)
#
## Dependencies
# - `gpg`
# - Python 3.7+

import base64
import click
import functools
import io
import json
import os
from os import path
import pexpect
import random
import stat
import string
import subprocess
import sys
import tempfile

## Secret loading/storing

def _load_store(ctx):
    secrets = ctx.find_object(dict)
    if secrets is not None:
        return secrets

    if path.exists('secrets.json.gpg'):
        process = subprocess.Popen(
            ['gpg', '--status-fd', '2', '--decrypt', '--quiet', 'secrets.json.gpg'],
            stdout = subprocess.PIPE,
            stderr = subprocess.PIPE,
            text = True,
        )

        if process.wait() != 0:
            raise RuntimeError('Failed to load secrets: %s', process.stderr.read())

        secrets = json.load(process.stdout)
    else:
        secrets = {"secrets": {}}

    ctx.obj = secrets

    return secrets

def _save_secrets(secrets):
    output = None
    try:
        output = tempfile.NamedTemporaryFile(dir = os.getcwd(), prefix = '.secrets-temp-', delete = False)

        process = subprocess.Popen(
            ['gpg', '--status-fd', '2', '--encrypt', '--quiet', '--no-tty', '--default-recipient-self'],
            stdin = subprocess.PIPE,
            stderr = subprocess.PIPE,
            stdout = output,
            text = True,
        )

        json.dump(secrets, process.stdin)
        process.stdin.close()

        if process.wait() != 0:
            raise click.ClickException(f'Failed to load secrets: {process.stderr.read()}')

        os.rename(output.name, 'secrets.json.gpg')
    finally:
        if output:
            try:
                os.unlink(output.name)
            except OSError:
                pass

def _migrate_secret(secret_type, secret, generate_kwargs):
    secret_version = secret["_secret_version"] if isinstance(secret, dict) else 0

    if secret_version < 1:
        secret = {
            "value": secret,
        }

    if secret_version < 2:
        secret["_generate_options"] = None

    if secret_version < 3:
        secret["_secret_type"] = secret_type

    secret["_secret_version"] = 3

    return secret

def _set_secret(store, secret_name, secret_type, secret_value, generate_kwargs):
    secret = store["secrets"][secret_name] = {
        "value": secret_value,
        "_secret_version": 3,
        "_secret_type": secret_type,
        "_generate_options": generate_kwargs,
    }

    return secret

def _run_secret_subcommand(ctx, secret_type, secret_generator, *, secret_name, generate, no_generate, **generate_kwargs):
    if no_generate:
        generate = 'no'

    store = _load_store(ctx)
    secret = store["secrets"].get(secret_name)

    if secret is None and generate == 'no':
        raise click.ClickException(f'Unknown secret: {secret_name}')

    if secret is not None:
        secret = _migrate_secret(secret_type, secret, generate_kwargs)

        if secret["_secret_type"] != secret_type:
            raise click.ClickException(f'Secret already exists as type "{secret["_secret_type"]}"')

    if secret is None or (generate == 'converge' and secret["_generate_options"] != generate_kwargs):
        secret = _set_secret(store, secret_name, secret_type, secret_generator(**generate_kwargs), generate_kwargs)
        _save_secrets(store)

    return secret["value"]

@click.group()
@click.pass_context
def _secretgarden(ctx):
    pass

def _secret_subcommand(func):
    @_secretgarden.command()
    @click.pass_context
    @click.argument('SECRET_NAME')
    @click.option(
        '--generate',
        type=click.Choice(('no', 'once', 'converge')),
        default='converge',
        show_default=True,
        help = 'Generate a secret if none exists. If set to `converge`, the default, the secret will be regenerated if the options have changed.',
    )
    @click.option(
        '-n',
        'no_generate',
        is_flag = True,
        help = 'Alias for `--generate=no`.',
    )
    @click.option(
        '--base64',
        'b64',
        is_flag=True,
        help = 'Encode result in base64',
    )
    @functools.wraps(func)
    def wrapper(ctx, *args, b64 = False, **kwargs):
        result = func(ctx, *args, **kwargs)

        if b64:
            result = base64.b64encode(result.encode('latin-1')).decode('latin-1')

        if getattr(ctx.find_root(), 'only_return_result', False):
            return result
        else:
            # SSH keys canonically have to contain a trailing newline, so we want to return it above
            # (but strip it for consistency here).
            print(result.rstrip())

def _generate_password(length):
    return ''.join(random.choice(string.ascii_letters + string.digits) for i in range(length))

@_secret_subcommand
@click.option(
    '--length',
    default=32,
    show_default=True,
    type=int,
    help = 'Generated password length',
)
def password(ctx, **kwargs):
    return _run_secret_subcommand(ctx, 'password', _generate_password, **kwargs)

def _generate_opaque():
    raise click.ClickException('Cannot generate opaque secret')

@_secret_subcommand
def opaque(ctx, **kwargs):
    return _run_secret_subcommand(ctx, 'opaque', _generate_opaque, **kwargs)

# `ssh-keygen` is awful to script. We force it to output the private key to a FIFO so it never lives
# on disk.
def _generate_ssh_key(key_type, key_bits = None):
    import pexpect
    with tempfile.TemporaryDirectory() as tmpdir:
        private_key_fifo_path = path.join(tmpdir, 'id')
        os.mkfifo(private_key_fifo_path)

        logfile = io.BytesIO()
        child = pexpect.spawn(
            'ssh-keygen',                    
            ['-q', '-N', '', '-t', key_type, '-f', private_key_fifo_path] + (['-b', str(key_bits)] if key_bits else []),
            timeout = 60,
            logfile = logfile,
        )

        try:
            child.expect(r'.+Overwrite \(y/n\)\? ')
            child.write('y\n')
            private_key = open(private_key_fifo_path).read()
            open(private_key_fifo_path + '.pub').read()
            child.read()
            child.close()
            child.wait()

            if child.exitstatus != 0:
                raise RuntimeError('Generating SSH key failed', logfile.getvalue())
        except KeyboardInterrupt as e:
            raise RuntimeError('Generating SSH key interrupted', logfile.getvalue()) from e
        except pexpect.TIMEOUT:
            raise RuntimeError('Generating SSH key timed out', logfile.getvalue()) from e

        return private_key

def _convert_ssh_private_key_to_public(private_key):
    import pexpect
    with tempfile.TemporaryDirectory() as tmpdir:
        private_key_fifo_path = path.join(tmpdir, 'id')
        os.mkfifo(private_key_fifo_path)
        os.chmod(private_key_fifo_path, stat.S_IRUSR | stat.S_IWUSR)

        logfile = io.StringIO()
        child = pexpect.spawn(
            'ssh-keygen',
            ['-y', '-f', private_key_fifo_path],
            timeout = 5,
            logfile = logfile,
            encoding = 'latin-1',
        )

        try:
            open(private_key_fifo_path, 'w', encoding = 'latin-1').write(private_key)
            public_key = child.read()
            child.close()
            child.wait()

            if child.exitstatus != 0:
                raise RuntimeError('Converting SSH key failed', logfile.getvalue())
        except KeyboardInterrupt as e:
            raise RuntimeError('Converting SSH key interrupted', logfile.getvalue()) from e
        except pexpect.TIMEOUT as e:
            raise RuntimeError('Converting SSH key timed out', logfile.getvalue()) from e

        return public_key

@_secret_subcommand
@click.option(
    '--public',
    is_flag=True,
    help = 'Output the public key, rather than the private key',
)
@click.option(
    '-t', '--type',
    'key_type',
    type = click.Choice(['dsa', 'ecdsa', 'ecdsa-sk', 'ed25519', 'ed25519-sk', 'rsa']),
    default = 'ed25519',
    show_default=True,
    help = 'Type of SSH key',
)
@click.option(
    '-b', '--bits',
    'key_bits',
    # Maximum number of bits supported by ssh-keygen.
    type = click.IntRange(min = 1, max = 16384),
    help = 'Type of SSH key',
)
def ssh_key(ctx, public, **kwargs):
    private_key = _run_secret_subcommand(ctx, 'ssh-key', _generate_ssh_key, **kwargs)

    if public:
        return _convert_ssh_private_key_to_public(private_key)
    else:
        return private_key

@_secretgarden.command()
@click.pass_context
@click.argument('SECRET_TYPE')
@click.argument('SECRET_NAME')
@click.argument('secret_value', metavar='VALUE', nargs = 1, required = False)
@click.option(
    '--base64',
    'b64',
    is_flag=True,
    help = 'Decode entered value as base64',
)
def set(ctx, *, secret_type, secret_name, secret_value, b64):
    if not secret_value:
        secret_value = sys.stdin.read().rstrip()

    if b64:
        secret_value = base64.b64decode(secret_value).decode('latin-1')

    store = _load_store(ctx)
    _set_secret(store, secret_name, secret_type, secret_value, {})
    _save_secrets(store)

def main():
    _secretgarden()

if __name__ == '__main__':
    main()
else:
    ## Ansible plugin
    DOCUMENTATION = """
        lookup: secretgarden
        author: Jesse Weaver <pianohacker@gmail.com>
        short_description: read/generate secrets
        description:
            - This lookup returns the contents from a file on the Ansible controller's file system.
        options:
          _terms:
            description: path(s) of files to read
            required: True
          rstrip:
            description: whether or not to remove whitespace from the ending of the looked-up file
            type: bool
            required: False
            default: True
          lstrip:
            description: whether or not to remove whitespace from the beginning of the looked-up file
            type: bool
            required: False
            default: False
        notes:
          - if read in variable context, the file can be interpreted as YAML if the content is valid to the parser.
          - this lookup does not understand 'globing', use the fileglob lookup instead.
    """

    from ansible.plugins.lookup import LookupBase
    from ansible.module_utils.six.moves import shlex_quote

    class LookupModule(LookupBase):
        def run(self, terms, variables, **kwargs):
            ctx = _secretgarden.make_context('secretgarden', terms + [
                (
                    (f'--{k}' if v else '')
                    if isinstance(v, bool) else
                    f'--{k}={shlex_quote(str(v))}' 
                ) for (k, v) in kwargs.items()
            ])
            ctx.only_return_result = True

            return [_secretgarden.invoke(ctx)]

# vim: set et :
