from cryptography.hazmat.primitives import serialization
import os
from cryptography.hazmat.primitives.asymmetric import ed25519
from pathlib import Path
import pyrage
from subprocess import Popen, PIPE, STDOUT


def gen_hostkey():
    private_key = ed25519.Ed25519PrivateKey.generate()
    public_key = private_key.public_key()

    private_bytes = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.OpenSSH,
        encryption_algorithm=serialization.NoEncryption())

    public_bytes = public_key.public_bytes(
        encoding=serialization.Encoding.OpenSSH,
        format=serialization.PublicFormat.OpenSSH)

    return private_bytes, public_bytes


def get_keys(key):
    private_key = serialization.load_ssh_private_key(key, password=None)
    public_key = private_key.public_key()

    private_bytes = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.OpenSSH,
        encryption_algorithm=serialization.NoEncryption())

    public_bytes = public_key.public_bytes(
        encoding=serialization.Encoding.OpenSSH,
        format=serialization.PublicFormat.OpenSSH)

    return private_bytes, public_bytes


admin_key = pyrage.ssh.Identity.from_buffer(
    open(os.path.expanduser("~/.ssh/id_ed25519"), "rb").read())
admin_pubkey = pyrage.ssh.Recipient.from_str(
    open(os.path.expanduser("~/.ssh/id_ed25519.pub"), "r").read())


def write_key(hostname, priv_key, pub_key, print_key=True):
    p = (Path('ssh-host-keys') / f"{hostname}.age")
    pub_path = (Path('ssh-host-keys') / f"{hostname}.pub")

    pub_path.write_text(pub_key.decode("utf-8"))
    p.write_bytes(pyrage.encrypt(priv_key, [admin_pubkey]))

    if print_key:
        print(pyrage.decrypt(p.read_bytes(), [admin_key]).decode("utf-8"))
    update_sops(hostname, pub_key)


def update_sops(hostname, pub_key):
    p = Popen(['ssh-to-age'], stdout=PIPE, stdin=PIPE, stderr=PIPE, text=True)
    stdout_data = p.communicate(input=pub_key.decode("utf-8"))[0]
    age_pubkey = stdout_data.strip()

    contents = open(".sops.yaml", "r").readlines()
    index = contents.index("keys:\n")

    found = False
    found_i = -1
    for i, line in enumerate(contents):
        if f"- &host_{hostname}" in line:
            found = True
            found_i = i
            break

    if not found:
        contents.insert(index + 1, f"- &host_{hostname} {age_pubkey}\n")
    else:
        contents[found_i] = f"- &host_{hostname} {age_pubkey}\n"

    open(".sops.yaml", "w").write(''.join(contents))