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


"""
TODO: this will be 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
- maybe one day: manage SSL certificates and keys

It uses files placed in <module_name>/secrets/* to know what secrets
it should handle. These secret files contain directives for what to do
about these secrets.

Example directives:

USER <description>
(a secret that must be filled in by the user)

USER_LONG <description>
(the same, indicates that the secret fits on several lines)

CMD <command>
(a secret that is generated by running this command)

CMD_ONCE <command>
(same, but value is not changed when doing a regen)

CONST <constant value>
(the secret has a constant value set here)

CONST_LONG
<constant value, several lines>
(same)

SERVICE_DN <service name> <service description>
(the LDAP DN of a service user)

SERVICE_PASSWORD <service name>
(the LDAP password for the corresponding service user)

SSL_CERT <cert name> <list of domains>
(a SSL domain for the given domains)

SSL_KEY <cert name>
(the SSL key going with corresponding certificate)

RSA_PUBLIC_KEY <key name> <key description>
(a public RSA key)

RSA_PRIVATE_KEY <key name>
(the corresponding private RSA key)
"""


# Parameters
LDAP_URL = "ldap://localhost:1389"
SERVICE_DN_SUFFIX = "ou=services,ou=users,dc=deuxfleurs,dc=fr"
consul_server = consul.Consul()


# ---- UTIL ----

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'

ldap_admin_conn = None
def get_ldap_admin_conn():
    global ldap_admin_conn
    if ldap_admin_conn is None:
        ldap_admin_conn = ldap.initialize(LDAP_URL)
        ldap_user = input("LDAP admin user (full DN, please!): ")
        ldap_pass = getpass.getpass("LDAP admin password: ")
        ldap_admin_conn.simple_bind_s(ldap_user, ldap_pass)
    return ldap_admin_conn

# ---- SECRETS ----

class Secret:
    def __init__(self, key, description=None):
        self.consul_key = "secrets/" + key
        if description != None:
            self.description = description

    def check(self, value):
        return True

    def generate(self, old_value):
        pass

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

    def generate(self, old_value):
            print("Description:", self.description)
            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):
    def __init__(self, command, rotate=False, **kwargs):
        Secret.__init__(self, **kwargs)
        self.command = command
        self.rotate = rotate

    def generate(self, old_value):
        if not self.rotate and old_value is not None:
            return None

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

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

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

    def generate(self, old_value):
        if old_value != self.value:
            return self.value
        else:
            return None

# ---- SERVICE USERS ----

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

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

    def generate(self, old_value):
        if old_value != self.service_user.password:
            return self.service_user.password
        else:
            return None

class ServiceUser:
    def __init__(self, username, password_secret, user_dn_secret=None, user_name_secret=None, rotate_password=False):
        self.username = username
        self.password = None
        self.dn = "cn={},{}".format(self.username, SERVICE_DN_SUFFIX)
        self.rotate_password = rotate_password

        self.username_secret = username_secret
        self.password_secret = password_secret
        self.dn_secret = dn_secret

    def secrets(self):
        secrets = {}
        secrets[self.password_secret] = ServiceUserPasswordSecret(
                service_user=self,
                key=self.password_secret,
                description="LDAP password for service user {}".format(self.username),
                )
        if self.user_dn_secret != None:
            secrets[self.dn_secret] = ConstantSecret(
                    key=self.dn_secret,
                    value=self.dn,
                    description="LDAP DN for service user {}".format(self.username),
                )
        if self.username_secret != None:
            secrets[self.username_secret] = ConstantSecret(
                    key=self.username_secret,
                    value=self.username,
                    description="LDAP username for service user {}".format(self.username),
                    )
        return secrets

# ---- READING secrets.toml FILES ----

class Config:
    def __init__(self, cluster_name):
        self.cluster_name = cluster_name
        self.app_path = os.path.join(".", "cluster", cluster_name)
        self.service_users = {}
        self.secrets = {}
        self.modules = []

    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, **uargs)
                for (skey, secret) in service_user.secrets():
                    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(key=skey, **sargs)
                elif ty == "command":
                    secret = CommandSecret(key=skey, **sargs)
                elif ty == "constant":
                    secret = Constantsecret(key=skey, **sargs)
                else:
                    raise Exception("Invalid secret type: {}".format(ty))
                if skey in self.secrets:
                    raise Exception("Duplicate secret: {}".format(skey))
                self.secrets[skey] = secret

    def check_secrets(self):
        print(":: Checking secrets...")
        for (_, secret) in self.secrets.items():
            _, data = consul_server.kv.get(secret.consul_key)
            if data is None:
                print(secret.consul_key, bcolors.FAIL, "x", bcolors.ENDC, " missing")
            elif not secret.check(data["Value"].decode('ascii').strip()): 
                print(secret.consul_key, bcolors.FAIL, "x", bcolors.ENDC, " wrong value")
            else:
                print(secret.consul_key, bcolors.OKGREEN, "✓", bcolors.ENDC)
        print()

    def gen_secrets(self):
        print(":: Generating missing secrets...")
        for (_, secret) in self.secrets.items():
            _, data = consul_server.kv.get(secret.consul_key)
            if data is None:
                value = secret.generate(None)
                if value != None:
                    print("PUT {} = {}".format(secret.consul_key, value))
                    print(bcolors.OKCYAN, "Value set.", bcolors.ENDC)
                else:
                    print(bcolors.WARNING, "Skipped.", bcolors.ENDC)
        print()

    def rotate_secrets(self):
        print(":: Rotating secrets...")
        for (_, secret) in self.secrets.items():
            print(": {} ({})".format(secret.consul_key, secret.__class__.__name__))
            _, data = consul_server.kv.get(secret.consul_key)
            if data is None:
                value = secret.generate(None)
            else:
                value = secret.generate(data["Value"].decode('ascii').strip())
            if value != None:
                print("PUT {} = {}".format(secret.consul_key, value))
                print(bcolors.OKCYAN, "Value set.", bcolors.ENDC)
            else:
                print(bcolors.WARNING, "Skipped.", bcolors.ENDC)
            print()

def check_secrets(modules):
    for mod in modules:
        print("==== MODULE `{}` ====".format(mod))
        modcfg = Config(mod)
        modcfg.check_secrets()
        print("")

def gen_secrets(modules):
    for mod in modules:
        print("==== MODULE `{}` ====".format(mod))
        modcfg = Config(mod)
        modcfg.gen_secrets()
        modcfg.check_secrets()
        print("")

def rotate_secrets(modules):
    for mod in modules:
        print("==== MODULE `{}` ====".format(mod))
        modcfg = Config(mod)
        modcfg.rotate_secrets()
        modcfg.check_secrets()
        print("")

# ----- old ----


# ---- CHECK COMMAND ----


def check_secrets_services(secrets):
    print("Checking secrets for LDAP service users...")
    services = get_secrets_services(secrets)

    for svc_name, svc in services.items():
        for dn_key in svc["dn_at"]:
            _, data = consul_server.kv.get(dn_key)
            if data is not None:
                got_val = data["Value"].decode('ascii').strip()
                if got_val != svc["dn"]:
                    print(svc_name, "wrong DN at", dn_key, bcolors.FAIL, "x", bcolors.ENDC)
                    print("got:", got_val, "instead of:", svc["dn"])

        if svc["pass"] is None:
            print(svc_name, bcolors.FAIL, "no password stored", bcolors.ENDC)
        else:
            for pass_key in svc["pass_at"]:
                _, data = consul_server.kv.get(pass_key)
                if data is not None:
                    got_val = data["Value"].decode('ascii').strip()
                    if got_val != svc["pass"]:
                        print(svc_name, "wrong pass at", dn_key, bcolors.FAIL, "x", bcolors.ENDC)

            l = ldap.initialize(LDAP_URL)
            try:
                l.simple_bind_s(svc["dn"], svc["pass"])
                print(svc_name, bcolors.OKGREEN, "✓", bcolors.ENDC)
            except Exception as e:
                print(svc_name, bcolors.FAIL, e, bcolors.ENDC)
    print()


# ---- GEN COMMAND ----

def gen_secrets_services(secrets, regen):
    print("Generating LDAP service accounts...")
    services = get_secrets_services(secrets)

    for svc_name, svc in services.items():
        print("----")
        print("Service:", svc_name)
        print("Description:", svc["desc"])

        for dn_key in svc["dn_at"]:
            _, data = consul_server.kv.get(dn_key)
            if data is None or data["Value"].decode('ascii').strip() != svc["dn"]:
                print(bcolors.OKCYAN, "Setting DN", bcolors.ENDC, "at", dn_key)
                consul_server.kv.put(dn_key, svc["dn"])

        if svc["pass"] is None or regen:
            print(bcolors.OKCYAN, "Generating new password", bcolors.ENDC)
            svc["pass"] = base64.urlsafe_b64encode(token_bytes(12)).decode('ascii')

        l = ldap.initialize(LDAP_URL)
        try:
            l.simple_bind_s(svc["dn"], svc["pass"])
        except:
            fix_service_user(svc)

        for pass_key in svc["pass_at"]:
            _, data = consul_server.kv.get(pass_key)
            if data is None or data["Value"].decode('ascii').strip() != svc["pass"]:
                print(bcolors.OKCYAN, "Setting password", bcolors.ENDC, "at", pass_key)
                consul_server.kv.put(pass_key, svc["pass"])

    print()

def fix_service_user(svc):
    print("Fixing service user", svc["dn"], "...")
    l = get_ldap_admin_conn()
    res = l.search_s(svc["dn"], ldap.SCOPE_BASE, "objectclass=*")
    pass_crypt = ldap_salted_sha1.hash(svc["pass"])
    if res is None or len(res) == 0:
        print(bcolors.OKCYAN, "Creating entity...", bcolors.ENDC)
        l.add_s(svc["dn"],
                [
                    ("objectclass",     [b"person", b"top"]),
                    ("displayname",     [svc["desc"].encode('ascii')]),
                    ("userpassword",     [pass_crypt.encode('ascii')]),
                ])
    else:
        print(bcolors.OKCYAN, "Resetting entity password", bcolors.ENDC)
        l.modify_s(svc["dn"],
                [
                    (ldap.MOD_REPLACE, "userpassword", [pass_crypt.encode('ascii')])
                ])

# ---- MAIN ----

if __name__ == "__main__":
    for i, val in enumerate(sys.argv):
        if val == "check":
            check_secrets(sys.argv[i+1:])
            break
        elif val == "gen":
            gen_secrets(sys.argv[i+1:])
            break
        elif val == "rotate":
            rotate_secrets(sys.argv[i+1:])
            break
        else:
            print("Usage:")
            print("    secretmgr.py [check|gen|rotate] <module name>...")


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