#!/usr/bin/env python

import subprocess
import sys
import typing as T

import yaml

ENCODING = "utf-8"
HOSTS: T.Dict[T.Text, T.Dict[T.Text, T.Text]] = {}


# From: https://stackoverflow.com/a/43060743
class DummyVault(yaml.YAMLObject):
    yaml_tag = "!vault"

    def __init__(self, cyphertext):
        self.cyphertext = "\n".join([e.strip() for e in cyphertext.split("\n")])

    def __repr__(self):
        return f"{self.__class__.__name__}({self.cyphertext[:10]}...)"

    @classmethod
    def from_yaml(cls, loader, node):
        return DummyVault(node.value)

    @classmethod
    def to_yaml(cls, dumper, data):
        return dumper.represent_scalar(cls.yaml_tag, data.cyphertext, style="|")


yaml.SafeLoader.add_constructor("!vault", DummyVault.from_yaml)
yaml.SafeDumper.add_multi_representer(DummyVault, DummyVault.to_yaml)


def load_hosts(inventory: T.Text) -> T.List[T.Text]:
    with open(inventory, "r") as f:
        data = yaml.load(f, Loader=yaml.SafeLoader)

    return [k for k in data["all"]["hosts"].keys()]


def gen_key(name: T.Text) -> None:
    with open(f"{name}.sec", "w") as sec:
        subprocess.call(["wg", "genkey"], stdout=sec)
    with open(f"{name}.sec", "rb", 0) as sec_r, open(f"{name}.pub", "w") as pub:
        subprocess.call(["wg", "pubkey"], stdin=sec_r, stdout=pub)


def to_vault(name: T.Text, passfile: T.Text) -> None:
    with open(f"{name}.pub", "r") as pub:
        pubkey = pub.readline()
        enc_pub = subprocess.check_output(
            [
                "ansible-vault",
                "encrypt_string",
                f"--vault-password-file={passfile}",
                pubkey,
            ]
        )
    with open(f"{name}.sec", "r") as sec:
        seckey = sec.readline()
        enc_sec = subprocess.check_output(
            [
                "ansible-vault",
                "encrypt_string",
                f"--vault-password-file={passfile}",
                seckey,
            ]
        )
    HOSTS[name] = {
        "public_key": yaml.load(enc_pub.decode(ENCODING), Loader=yaml.SafeLoader),
        "private_key": yaml.load(enc_sec.decode(ENCODING), Loader=yaml.SafeLoader),
    }


def usage() -> None:
    print("Usage: \n\tgen_keys.py <path/to/inventory> <path/to/vault_password_file>")


if __name__ == "__main__":
    if len(sys.argv) != 3:
        usage()
        sys.exit(-1)

    for host in load_hosts(sys.argv[1]):
        gen_key(host)
        to_vault(host, sys.argv[2])
        result = yaml.dump(HOSTS, Dumper=yaml.SafeDumper)
        with open("result.yml", "w") as res:
            res.write(result)
        print(result)