#!/usr/bin/env python3
"""
pinentry-vault -- PIN-protected credential store for pinentry.

Store passphrases encrypted with a short PIN.  When used as a pinentry
wrapper, gpg-agent calls pinentry-vault; the user enters a PIN instead
of the real passphrase, and pinentry-vault decrypts and returns the
actual passphrase.

USAGE:
  pinentry-vault --ls-key                     List all keys
  pinentry-vault <key> --delete               Delete a key
  pinentry-vault <key> --add <passwd> <pin>   Create/update a key
  pinentry-vault <key> [pinentry-args ...]    Pinentry wrapper mode

In wrapper mode the script is meant to be set as pinentry-program in
gpg-agent.conf:

  pinentry-program /home/USER/.local/bin/pinentry-vault mykey
"""

import argparse
import os
import subprocess
import sys


VAULT_DIR = os.path.expanduser("~/.local/share/pinentry-vault")
DEFAULT_PINENTRY = "pinentry"


# ---------------------------------------------------------------------------
#  Percent encoding / decoding  (Assuan protocol)
# ---------------------------------------------------------------------------

def percent_encode(s: str) -> str:
    """Percent-encode a string for use in an Assuan D-line."""
    out = []
    for ch in s:
        b = ord(ch)
        if ch == "%":
            out.append("%25")
        elif b < 0x20 or b > 0x7E:
            out.append("%{:02X}".format(b))
        else:
            out.append(ch)
    return "".join(out)


def percent_decode(s: str) -> str:
    """Decode an Assuan percent-encoded string."""
    out = []
    i = 0
    n = len(s)
    while i < n:
        if s[i] == "%" and i + 2 < n:
            try:
                out.append(chr(int(s[i + 1 : i + 3], 16)))
                i += 3
                continue
            except ValueError:
                pass
        out.append(s[i])
        i += 1
    return "".join(out)


# ---------------------------------------------------------------------------
#  Encryption / Decryption  (openssl)
# ---------------------------------------------------------------------------

def _run_openssl(args, input_data, pin):
    """Run openssl subprocess, passing PIN via environment variable."""
    env = os.environ.copy()
    env["VAULT_PIN"] = pin
    proc = subprocess.run(
        ["openssl"] + args,
        input=input_data.encode(),
        capture_output=True,
        env=env,
    )
    if proc.returncode != 0:
        return None
    return proc.stdout.decode()


_ENCRYPT_ARGS = ["enc", "-aes-256-cbc", "-pbkdf2", "-iter", "100000",
                 "-salt", "-a", "-pass", "env:VAULT_PIN"]
_DECRYPT_ARGS = ["enc", "-d", "-aes-256-cbc", "-pbkdf2", "-iter", "100000",
                 "-salt", "-a", "-pass", "env:VAULT_PIN"]


def encrypt(passwd: str, pin: str) -> str:
    """Encrypt *passwd* with *pin*, return base64-encoded ciphertext."""
    result = _run_openssl(_ENCRYPT_ARGS, passwd, pin)
    if result is None:
        sys.exit("error: encryption failed")
    return result


def decrypt(ciphertext: str, pin: str):
    """Decrypt *ciphertext* with *pin*, return passwd or None on failure."""
    return _run_openssl(_DECRYPT_ARGS, ciphertext, pin)


# ---------------------------------------------------------------------------
#  Key storage helpers
# ---------------------------------------------------------------------------

def key_path(key: str) -> str:
    return os.path.join(VAULT_DIR, key + ".enc")


def ensure_vault_dir():
    os.makedirs(VAULT_DIR, mode=0o700, exist_ok=True)


# ---------------------------------------------------------------------------
#  Commands
# ---------------------------------------------------------------------------

def cmd_ls_key():
    if not os.path.isdir(VAULT_DIR):
        return
    for name in sorted(os.listdir(VAULT_DIR)):
        if name.endswith(".enc"):
            print(name[:-4])


def cmd_delete(key: str):
    path = key_path(key)
    if os.path.isfile(path):
        os.remove(path)
        print("Deleted:", key)
    else:
        print("error: key not found:", key, file=sys.stderr)
        sys.exit(1)


def cmd_add(key: str, passwd: str, pin: str):
    ensure_vault_dir()
    encrypted = encrypt(passwd, pin)
    path = key_path(key)
    old_umask = os.umask(0o077)
    try:
        with open(path, "w") as f:
            f.write(encrypted)
    finally:
        os.umask(old_umask)
    print("Added:", key)


def cmd_wrapper(key: str, pinentry_args):
    """Pinentry wrapper mode: forward protocol, intercept GETPIN."""
    path = key_path(key)
    if not os.path.isfile(path):
        _assuan_err("83886179 No vault entry for '%s'" % key)
        sys.exit(0)

    with open(path) as f:
        encrypted = f.read()

    pinentry_bin = os.environ.get("VAULT_PINENTRY", DEFAULT_PINENTRY)
    proc = subprocess.Popen(
        [pinentry_bin] + pinentry_args,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=sys.stderr,
        text=True,
        bufsize=1,
    )

    try:
        _forward_greeting(proc)

        for line in sys.stdin:
            cmd = line.strip()
            if not cmd:
                continue

            parts = cmd.split(None, 1)
            cmd_name = parts[0].upper() if parts else ""

            if cmd_name == "GETPIN":
                _handle_getpin(proc, key, encrypted)
            elif cmd_name == "SETDESC":
                _forward_line(proc, cmd + " [vault: %s]" % key)
                _forward_response(proc)
            elif cmd_name == "BYE":
                _forward_line(proc, "BYE")
                _forward_response(proc)
                break
            else:
                _forward_line(proc, cmd)
                _forward_response(proc)
    except (BrokenPipeError, KeyboardInterrupt):
        pass
    finally:
        _kill_proc(proc)


# ---------------------------------------------------------------------------
#  Pinentry protocol helpers
# ---------------------------------------------------------------------------

def _forward_greeting(proc):
    """Read and forward pinentry's initial greeting line."""
    greeting = proc.stdout.readline()
    if greeting:
        sys.stdout.write(greeting)
        sys.stdout.flush()


def _forward_line(proc, text):
    """Send a single line to pinentry."""
    proc.stdin.write(text + "\n")
    proc.stdin.flush()


def _forward_response(proc):
    """Read pinentry response lines and forward them to stdout."""
    while True:
        line = proc.stdout.readline()
        if not line:
            break
        line = line.rstrip("\n")
        if line:
            sys.stdout.write(line + "\n")
            sys.stdout.flush()
        if line == "OK" or line.startswith("ERR"):
            break


def _handle_getpin(proc, key, encrypted):
    """Intercept GETPIN.

    Forward GETPIN to the real pinentry so the user can type their PIN.
    Then decrypt the vault entry with that PIN and return the real passwd.
    """
    _forward_line(proc, "GETPIN")

    pin = None
    status_lines = []

    while True:
        line = proc.stdout.readline()
        if not line:
            return
        line = line.rstrip("\n")

        if line.startswith("S "):
            status_lines.append(line)
        elif line.startswith("D "):
            pin = percent_decode(line[2:])
        elif line == "OK":
            break
        elif line.startswith("ERR"):
            sys.stdout.write(line + "\n")
            sys.stdout.flush()
            return

    if pin is None or pin == "":
        _assuan_err("83886179 No PIN entered")
        return

    passwd = decrypt(encrypted, pin)
    if passwd is None:
        _assuan_err("83886179 Bad PIN")
        return

    for s in status_lines:
        sys.stdout.write(s + "\n")

    sys.stdout.write("D " + percent_encode(passwd) + "\n")
    sys.stdout.write("OK\n")
    sys.stdout.flush()


def _assuan_err(msg):
    """Write an Assuan ERR response."""
    sys.stdout.write("ERR %s <Pinentry>\n" % msg)
    sys.stdout.flush()


def _kill_proc(proc):
    """Terminate pinentry subprocess cleanly."""
    try:
        proc.terminate()
        proc.wait(timeout=3)
    except Exception:
        try:
            proc.kill()
            proc.wait(timeout=3)
        except Exception:
            pass


# ---------------------------------------------------------------------------
#  Argument parsing
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="PIN-protected credential store for pinentry",
        add_help=False,
    )
    parser.add_argument("--help", action="store_true")
    parser.add_argument("--ls-key", action="store_true")
    parser.add_argument("--delete", action="store_true")
    parser.add_argument("--add", nargs=2, metavar=("PASSWD", "PIN"))
    parser.add_argument("key", nargs="?")

    args, unknown = parser.parse_known_args()

    if args.help:
        print(__doc__.strip())
        return

    if args.ls_key:
        cmd_ls_key()
        return

    if args.delete:
        if not args.key:
            sys.exit("error: KEY required for --delete")
        cmd_delete(args.key)
        return

    if args.add:
        if not args.key:
            sys.exit("error: KEY required for --add")
        passwd, pin = args.add
        cmd_add(args.key, passwd, pin)
        return

    if args.key:
        cmd_wrapper(args.key, unknown)
    else:
        print(__doc__.strip(), file=sys.stderr)


if __name__ == "__main__":
    main()