#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p "python3.withPackages(ps: [ ps.pip ps.consul ps.ldap ps.passlib ps.requests ps.six ps.toml ])"

# DEPENDENCY: python-consul
import consul

# DEPENDENCY: python-ldap
import ldap

# DEPENDENCY: passlib
from passlib.hash import ldap_salted_sha1

# DEPENDENCY: toml
import toml

import os
import sys
import glob
import subprocess
import getpass
import base64
from secrets import token_bytes

"""
This is a utility to handle secrets in the Consul database
for the various components of the Deuxfleurs infrastructure

Functionnalities:
- check that secrets are correctly configured
- help user fill in secrets
- create LDAP service users and fill in corresponding secrets
- TODO: manage Garage buckets and access keys
- maybe one day: manage SSL certificates and keys
"""

# ---- UTIL ----

consul_server = consul.Consul()

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


# ---- SECRETS ----

class Secret:
    def __init__(self, key, config, description=None):
        self.config = config
        self.key = key
        if description != None:
            self.description = description
        else:
            self.description = None

    def check(self, value):
        return True

    def generate(self):
        pass

    def rotate(self):
        return None

    def print_info(self):
        print("Secret:         {}".format(self.key))
        print("Type:           {}".format(self.__class__.TYPE))
        if self.description != None:
            print("Description:    {}".format(self.description))

class UserSecret(Secret):
    TYPE = "user-entered secret"

    def __init__(self, example=None, multiline=False, **kwargs):
        Secret.__init__(self, **kwargs)
        self.example = example
        self.multiline = multiline

    def print_info(self):
        Secret.print_info(self)
        if self.key in self.config.user_values:
            print("Cluster value:  {}".format(self.config.user_values[self.key]))
        elif self.example != None:
            print("Example:        {}".format(self.example))

    def generate(self):
        if self.key in self.config.user_values:
            print("Using constant value from cluster's secretmgr.toml")
            return self.config.user_values[self.key]

        print("Enter value for secret, or ^C to skip:")
        if self.multiline:
            print("THIS IS A LONG VALUE, ENTER SEVERAL LINES AND FINISH WITH A LINE CONTAINING A SINGLE .")
            try:
                lines = []
                while True:
                    line = input().strip()
                    if line == ".":
                        break
                    lines.append(line)
                return "\n".join(lines)
            except KeyboardInterrupt:
                return None
        else:
            try:
                while True:
                    line = input().strip()
                    if line != "":
                        return line
                    else:
                        print("Please enter a non-empty value, or ^C to skip:")
            except KeyboardInterrupt:
                return None

class CommandSecret(Secret):
    TYPE = "command"

    def __init__(self, command, rotate=False, **kwargs):
        Secret.__init__(self, **kwargs)
        self.command = command
        self.rotate_value = rotate

    def print_info(self):
        Secret.print_info(self)
        print("Command:        {}".format(self.command))
        if self.rotate_value:
            print("Rotate:         True")

    def generate(self):
        print("Executing command:", self.command)
        return subprocess.check_output(["sh", "-c", self.command])

    def rotate(self):
        if self.rotate_value:
            return self.generate()
        else:
            return None

class ConstantSecret(Secret):
    TYPE = "constant value"

    def __init__(self, value, **kwargs):
        Secret.__init__(self, **kwargs)
        self.value = value

    def print_info(self):
        Secret.print_info(self)
        print("Value:          {}".format(self.value))

    def check(self, value):
        return value == self.value

    def generate(self):
        return self.value


# ---- SERVICE USERS ----

class ServiceUserPasswordSecret(Secret):
    TYPE = "service user's password"

    def __init__(self, service_user, **kwargs):
        Secret.__init__(self, **kwargs)
        self.service_user = service_user

    def print_info(self):
        Secret.print_info(self)
        print("Service user:   {}".format(self.service_user.username))

    def check(self, value):
        l = ldap.initialize(self.config.ldap_server)
        try:
            l.simple_bind_s(self.service_user.dn, value)
            return True
        except Exception as e:
            return False

    def generate(self):
        return self.service_user.password

    def rotate(self):
        return self.service_user.password

class ServiceUserNameSecret(Secret):
    TYPE = "service user's username (constant value)"

    def __init__(self, service_user, **kwargs):
        Secret.__init__(self, **kwargs)
        self.service_user = service_user

    def print_info(self):
        Secret.print_info(self)
        print("Value:          {}".format(self.service_user.username))

    def check(self, value):
        return value == self.service_user.username

    def generate(self):
        return self.service_user.username

class ServiceUserDNSecret(Secret):
    TYPE = "service user's DN (constant value)"

    def __init__(self, service_user, **kwargs):
        Secret.__init__(self, **kwargs)
        self.service_user = service_user

    def print_info(self):
        Secret.print_info(self)
        print("Service user:   {}".format(self.service_user.username))
        print("Value:          {}".format(self.service_user.dn))

    def check(self, value):
        return value == self.service_user.dn

    def generate(self):
        return self.service_user.dn

class ServiceUser:
    def __init__(self, username, password_secret, config, description=None, dn_secret=None, username_secret=None, rotate_password=False):
        self.config = config
        self.username = username
        self.description = description
        self.password = None
        self.dn = "cn={},{}".format(self.username, self.config.ldap_service_dn_suffix)
        self.rotate_password = rotate_password

        self.password_secret = ServiceUserPasswordSecret(config=config, service_user=self, key=password_secret)

        self.username_secret = None
        if username_secret != None:
            self.username_secret = ServiceUserNameSecret(config=config, service_user=self, key=username_secret)

        self.dn_secret = None
        if dn_secret != None:
            self.dn_secret = ServiceUserDNSecret(config=config, service_user=self, key=dn_secret)

    def secrets(self):
        secrets = {}
        secrets[self.password_secret.key] = self.password_secret
        if self.dn_secret != None:
            secrets[self.dn_secret.key] = self.dn_secret
        if self.username_secret != None:
            secrets[self.username_secret.key] = self.username_secret
        return secrets

    def configure(self, rotate):
        self.password = self.config.get_secret(self.password_secret.key)
        if self.password is None:
            good = False
        else:
            l = ldap.initialize(self.config.ldap_server)
            try:
                l.simple_bind_s(self.dn, self.password)
                good = True
            except:
                good = False

        if not good or (rotate and self.rotate_password):
            # Reset passsword
            self.password = base64.urlsafe_b64encode(token_bytes(12)).decode('ascii')
            pass_crypt = ldap_salted_sha1.hash(self.password).encode('ascii')

            l = self.config.get_ldap_admin_conn()
            res = l.search_s(self.dn, ldap.SCOPE_BASE, "objectclass=*")
            if res is None or len(res) == 0:
                print(bcolors.OKCYAN, "Creating entity", self.dn, bcolors.ENDC)
                if self.config.dry_run:
                    print(bcolors.OKBLUE, "Dry run, skipping. Add --do to actually do something.", bcolors.ENDC)
                    return
                l.add_s(self.dn,
                        [
                            ("objectclass",     [b"person", b"top"]),
                            ("displayname",     [self.description.encode('ascii')]),
                            ("userpassword",     [pass_crypt]),
                        ])
            else:
                print(bcolors.OKCYAN, "Resetting password for entity", self.dn, bcolors.ENDC)
                if self.config.dry_run:
                    print(bcolors.OKBLUE, "Dry run, skipping. Add --do to actually do something.", bcolors.ENDC)
                    return
                l.modify_s(self.dn,
                        [
                            (ldap.MOD_REPLACE, "userpassword", [pass_crypt])
                        ])
        else:
            print(bcolors.OKGREEN, "Entity is good: ", self.dn, bcolors.ENDC)


# ---- MAIN CONFIG CLASS ----

class Config:
    def __init__(self, cluster_name, dry_run):
        self.cluster_name = cluster_name
        self.app_path = os.path.join(".", "cluster", cluster_name, "app")

        self.service_users = {}
        self.secrets = {}
        self.modules = []
        self.dry_run = dry_run

        # Load config from secretmgr.toml in cluster directory
        secretmgr_toml_path = os.path.join(".", "cluster", cluster_name, "secretmgr.toml")
        if os.path.exists(secretmgr_toml_path):
            with open(secretmgr_toml_path) as f:
                secretmgr_toml = toml.load(f)
        else:
            secretmgr_toml = {}

        if "user_values" in secretmgr_toml:
            self.user_values = secretmgr_toml["user_values"]
        else:
            self.user_values = {}

        self.ldap_server = None
        self.ldap_service_dn_suffix = None
        self.ldap_admin_dn = None
        self.ldap_admin_password_secret = None
        self.ldap_admin_conn = None
        if "ldap" in secretmgr_toml:
            if "server" in secretmgr_toml["ldap"]:
                self.ldap_server = secretmgr_toml["ldap"]["server"]
            if "service_dn_suffix" in secretmgr_toml["ldap"]:
                self.ldap_service_dn_suffix = secretmgr_toml["ldap"]["service_dn_suffix"]
            if "admin_dn" in secretmgr_toml["ldap"]:
                self.ldap_admin_dn = secretmgr_toml["ldap"]["admin_dn"]
            if "admin_password_secret" in secretmgr_toml["ldap"]:
                self.ldap_admin_password_secret = secretmgr_toml["ldap"]["admin_password_secret"]

    def load_module(self, module_name):
        secrets_toml_path = os.path.join(self.app_path, module_name, "secrets.toml")

        with open(secrets_toml_path) as f:
            secrets_toml = toml.load(f)

        self.modules.append(module_name)

        # Service users, and their associated secrets
        if "service_users" in secrets_toml:
            for (uname, uargs) in secrets_toml["service_users"].items():
                service_user = ServiceUser(uname, config=self, **uargs)
                for (skey, secret) in service_user.secrets().items():
                    if skey in self.secrets:
                        raise Exception("Duplicate secret: {}".format(skey))
                    self.secrets[skey] = secret
                self.service_users[uname] = service_user

        # Other secrets
        if "secrets" in secrets_toml:
            for (skey, sargs) in secrets_toml["secrets"].items():
                ty = sargs["type"]
                del sargs["type"]
                if ty == "user":
                    secret = UserSecret(config=self, key=skey, **sargs)
                elif ty == "command":
                    secret = CommandSecret(config=self, key=skey, **sargs)
                elif ty == "constant":
                    secret = ConstantSecret(config=self, key=skey, **sargs)
                elif ty == "service_password":
                    service = sargs["service"]
                    del sargs["service"]
                    secret = ServiceUserPasswordSecret(
                        config=self,
                        key=skey,
                        service_user=self.service_users[service],
                        **sargs)
                elif ty == "service_username":
                    service = sargs["service"]
                    del sargs["service"]
                    secret = ServiceUserNameSecret(
                        config=self,
                        key=skey,
                        service_user=self.service_users[service],
                        **sargs)
                elif ty == "service_dn":
                    service = sargs["service"]
                    del sargs["service"]
                    secret = ServiceUserDNSecret(
                        config=self,
                        key=skey,
                        service_user=self.service_users[service],
                        **sargs)
                else:
                    description = "{}, {}".format(ty,
                                ", ".join([k + ": " + v for k, v in sargs.items()]))
                    secret = UserSecret(
                        config=self,
                        key=skey,
                        multiline=True,
                        description=description)
                if skey in self.secrets:
                    raise Exception("Duplicate secret: {}".format(skey))
                self.secrets[skey] = secret

    def add_user_values_secrets(self):
        for (skey, value) in self.user_values.items():
            self.secrets[skey] = ConstantSecret(
                config=self,
                key=skey,
                value=value,
                description="Cluster-defined user value")

    # -- consul and ldap helpers --

    def check_consul_cluster(self):
        # Check cluster name we are connected to
        consul_node = consul_server.agent.self()
        if consul_node["Config"]["Datacenter"] != self.cluster_name:
            print("You are not connected to the correct Consul cluster.")
            print("You are connected to cluster '{}' instead of '{}'.".format(consul_node["Config"]["Datacenter"], self.cluster_name))
            sys.exit(1)

    def get_ldap_admin_conn(self):
        if self.ldap_admin_conn is None:
            if self.ldap_admin_password_secret != None:
                ldap_pass = self.get_secret(self.ldap_admin_password_secret)
                if ldap_pass is None:
                    raise Exception("LDAP admin password could not be read at: {}".format(pass_key))
            else:
                ldap_pass = getpass.getpass("LDAP admin password: ")

            self.ldap_admin_conn = ldap.initialize(self.ldap_server)
            self.ldap_admin_conn.simple_bind_s(self.ldap_admin_dn, ldap_pass)
        return self.ldap_admin_conn

    def get_secret(self, key):
        _, data = consul_server.kv.get("secrets/" + key)
        if data is None:
            return None
        else:
            return data["Value"].decode('ascii').strip()

    def put_secret(self, key, value):
        if self.dry_run:
            print(bcolors.OKBLUE, "Dry run, not updating secrets/{}. Add --do to actually do something.".format(key), bcolors.ENDC)
            return
        consul_server.kv.put("secrets/" + key, value)

    # -- user actions --

    def print_info(self):
        print("== LIST OF SERVICE USERS ==")
        print()
        for (_, su) in self.service_users.items():
            print("Username:      {}".format(su.username))
            print("DN:            {}".format(su.dn))
            print("Pass. secret:  {}".format(su.password_secret.key))
            print()

        print("== LIST OF SECRETS ==")
        print()
        for (_, secret) in self.secrets.items():
            secret.print_info()
            print()

    def check_secrets(self):
        self.check_consul_cluster()
        print(":: Checking secrets...")
        must_gen = False
        for (_, secret) in self.secrets.items():
            value = self.get_secret(secret.key)
            if value is None:
                print(secret.key, bcolors.FAIL, "x  missing", bcolors.ENDC)
                must_gen = True
            elif not secret.check(value): 
                print(secret.key, bcolors.WARNING, "x  bad value", bcolors.ENDC)
                must_gen = True
            else:
                print(secret.key, bcolors.OKGREEN, "✓", bcolors.ENDC)
        print()
        if must_gen:
            print("To fix missing or invalid secrets, use `secretmgr gen <cluster_name> <app>...`")
            print()

    def gen_secrets(self):
        self.check_consul_cluster()
        if len(self.service_users) > 0:
            print(":: Configuring service users...")
            for (_, su) in self.service_users.items():
                su.configure(False)
            print()

        print(":: Generating missing/invalid secrets...")
        for (_, secret) in self.secrets.items():
            old_value = self.get_secret(secret.key)
            if old_value is None or not secret.check(old_value):
                print()
                secret.print_info()
                value = secret.generate()
                if value != None:
                    self.put_secret(secret.key, value)
                    print(bcolors.OKCYAN, "Value set.", bcolors.ENDC)
                else:
                    print(bcolors.WARNING, "Skipped.", bcolors.ENDC)

        print()
        self.check_secrets()

    def rotate_secrets(self):
        self.check_consul_cluster()
        if len(self.service_users) > 0:
            print(":: Regenerating service user passwords...")
            for (_, su) in self.service_users.items():
                su.configure(True)
            print()

        print(":: Rotating secrets...")
        for (_, secret) in self.secrets.items():
            print()
            secret.print_info()

            old_value = self.get_secret(secret.key)
            new_value = secret.rotate()

            if new_value != None and new_value != old_value:
                self.put_secret(secret.key, new_value)
                print(bcolors.OKCYAN, "Value set.", bcolors.ENDC)
            else:
                print(bcolors.OKGREEN, "Nothing to do.", bcolors.ENDC)

        print()
        self.check_secrets()


# ---- MAIN ----

def load_config(cluster_name, modules, **kwargs):
    # Load config
    cfg = Config(cluster_name, **kwargs)
    if len(modules) > 0:
        for mod in modules:
            cfg.load_module(mod)
    else:
        cfg.add_user_values_secrets()

    return cfg

if __name__ == "__main__":
    verb = None
    dry_run = True

    for i, val in enumerate(sys.argv):
        if val == "--do":
            dry_run = False
        elif val == "info":
            verb = lambda cfg: cfg.print_info()
            break
        elif val == "check":
            verb = lambda cfg: cfg.check_secrets()
            break
        elif val == "gen":
            verb = lambda cfg: cfg.gen_secrets()
            break
        elif val == "rotate":
            verb = lambda cfg: cfg.rotate_secrets()
            break

    if verb is None:
        print("Usage:")
        print("    secretmgr [--do] info|check|gen|rotate <cluster name> [<module name>...]")
    else:
        cfg = load_config(
                cluster_name=sys.argv[i+1],
                modules=sys.argv[i+2:],
                dry_run=dry_run)
        verb(cfg)


#  vim: set sts=4 ts=4 sw=4 tw=0 ft=python et :