#!/usr/bin/env python3
#
# Copyright (c) 2024 Dhruva Sambrani <dhruvasambrani19@gmail.com>
#
# This file is part of the pam_usb project. pam_usb is free software;
# you can redistribute it and/or modify it under the terms of the GNU General
# Public License version 2, as published by the Free Software Foundation.
#
# pam_usb is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 51 Franklin
# Street, Fifth Floor, Boston, MA 02110-1301 USA.

import os
import sys
import subprocess
import getpass
import argparse
import shutil
import syslog
import pwd
from dotenv import load_dotenv


def _user_home():
    """Return the real user's home directory, even when invoked via sudo."""
    sudo_user = os.environ.get("SUDO_USER")
    if sudo_user and os.geteuid() == 0:
        return pwd.getpwnam(sudo_user).pw_dir
    return os.path.expanduser("~")


ENVFILE_PATH = os.path.join(_user_home(), ".pamusb", ".pinentry.env")
PINENTRY_LINK = "/usr/bin/pinentry"
PINENTRY_ALT_NAME = "pinentry"
PAMUSB_PINENTRY_PATH = "/usr/bin/pinentry-pamusb"
PAMUSB_CHECK_PATH = "/usr/bin/pamusb-check"
PINENTRY_ALT_PRIORITY = 100

ENVFILE_TEMPLATE = (
    "# pinentry-pamusb configuration\n"
    "# Your GPG passphrase for auto-unlock\n"
    "PINENTRY_PASSWORD=changeme\n"
    "\n"
    "# Fallback pinentry application (must be an absolute path)\n"
    "# PINENTRY_FALLBACK_APP=/usr/bin/pinentry-gnome3\n"
)


def install():
    envfile_dir = os.path.dirname(ENVFILE_PATH)
    os.makedirs(envfile_dir, exist_ok=True)

    if not os.path.exists(ENVFILE_PATH):
        with open(ENVFILE_PATH, 'w') as f:
            f.write(ENVFILE_TEMPLATE)
        os.chmod(ENVFILE_PATH, 0o600)
        print("Created %s" % ENVFILE_PATH)
    else:
        print("%s already exists, skipping creation." % ENVFILE_PATH)

    if shutil.which("update-alternatives") is None:
        print("update-alternatives not found, skipping alternative registration.")
        return

    result = subprocess.run(
        ["update-alternatives", "--install", PINENTRY_LINK, PINENTRY_ALT_NAME,
         PAMUSB_PINENTRY_PATH, str(PINENTRY_ALT_PRIORITY)],
        capture_output=True
    )
    if result.returncode != 0:
        sys.stderr.write("Failed to register alternative: %s\n" % result.stderr.decode())
        sys.exit(1)

    result = subprocess.run(
        ["update-alternatives", "--set", PINENTRY_ALT_NAME, PAMUSB_PINENTRY_PATH],
        capture_output=True
    )
    if result.returncode != 0:
        sys.stderr.write("Failed to activate alternative: %s\n" % result.stderr.decode())
        sys.exit(1)

    print("pinentry-pamusb registered and activated as pinentry alternative.")


def uninstall():
    if os.path.exists(ENVFILE_PATH):
        os.remove(ENVFILE_PATH)
        print("Removed %s" % ENVFILE_PATH)
    else:
        print("%s does not exist, skipping removal." % ENVFILE_PATH)

    if shutil.which("update-alternatives") is None:
        print("update-alternatives not found, skipping alternative removal.")
        return

    result = subprocess.run(
        ["update-alternatives", "--remove", PINENTRY_ALT_NAME, PAMUSB_PINENTRY_PATH],
        capture_output=True
    )
    if result.returncode != 0:
        sys.stderr.write("Failed to remove alternative: %s\n" % result.stderr.decode())
        sys.exit(1)

    print("pinentry-pamusb alternative removed.")

def escape_assuan_data(data):
    """Assuan protocol requires % and newlines to be percent-escaped in D lines."""
    if not data:
        return ""
    return data.replace('%', '%25').replace('\n', '%0A').replace('\r', '%0D')

def _run_pamusb_check(user):
    return subprocess.run([PAMUSB_CHECK_PATH, user], capture_output=True)


def _run_as_pinentry(user, pinentry_password, fallback_pinentry_app, unknown_args):
    syslog.openlog('pinentry-pamusb', syslog.LOG_PID, syslog.LOG_AUTH)

    if not pinentry_password:
        syslog.syslog(syslog.LOG_WARNING,
                      f"PINENTRY_PASSWORD is not set for user '{user}'")

    try:
        auth_result = _run_pamusb_check(user)
    except OSError as e:
        syslog.syslog(syslog.LOG_ERR,
                      f"Failed to run pamusb-check for user '{user}': {e}")
        auth_result = None

    if auth_result is not None and auth_result.returncode == 0:
        syslog.syslog(syslog.LOG_NOTICE,
                      f"Authentication succeeded for user '{user}'")
        print("OK Pleased to meet you", flush=True)
        while True:
            try:
                line_str = input().strip()
            except EOFError:
                break

            if not line_str or line_str.startswith('#'):
                continue

            cmd = line_str.split(maxsplit=1)[0]

            if cmd == "GETPIN":
                syslog.syslog(syslog.LOG_NOTICE,
                              f"Passphrase delivered to gpg-agent for user '{user}'")
                print("D %s" % escape_assuan_data(pinentry_password), flush=True)
                print("OK", flush=True)
            elif cmd == "BYE":
                print("OK", flush=True)
                break
            else:
                print("OK", flush=True)
    else:
        if auth_result is not None:
            reason = "; ".join(auth_result.stderr.decode(errors='replace').splitlines()).strip() or "unknown reason"
            syslog.syslog(syslog.LOG_NOTICE,
                          f"Authentication failed for user '{user}' (code {auth_result.returncode}): {reason}")
        if (fallback_pinentry_app
                and os.path.isabs(fallback_pinentry_app)
                and os.path.isfile(fallback_pinentry_app)
                and os.access(fallback_pinentry_app, os.X_OK)):
            syslog.syslog(syslog.LOG_NOTICE,
                          f"Falling back to '{fallback_pinentry_app}'")
            try:
                os.execv(fallback_pinentry_app, [fallback_pinentry_app] + unknown_args)
            except OSError as e:
                syslog.syslog(syslog.LOG_ERR,
                              f"Failed to execute fallback '{fallback_pinentry_app}': {e}")
                sys.stderr.write(f"Failed to execute fallback '{fallback_pinentry_app}': {e}\n")
                sys.exit(1)
        else:
            syslog.syslog(syslog.LOG_ERR,
                          f"PINENTRY_FALLBACK_APP='{fallback_pinentry_app}' is not set or not a valid executable, cannot fall back")
            sys.stderr.write(
                f'PINENTRY_FALLBACK_APP={fallback_pinentry_app} is not set or not a valid '
                f'executable path, cannot fall back.\n')
            sys.exit(1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="pam_usb pinentry tool"
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "--install", action="store_true",
        help="Create envfile and register pinentry-pamusb as pinentry alternative"
    )
    group.add_argument(
        "--uninstall", action="store_true",
        help="Remove envfile and unregister pinentry-pamusb pinentry alternative"
    )

    args, unknown_args = parser.parse_known_args()

    if args.install:
        install()
        sys.exit(0)

    if args.uninstall:
        uninstall()
        sys.exit(0)

    load_dotenv(ENVFILE_PATH)
    _run_as_pinentry(
        getpass.getuser(),
        os.getenv('PINENTRY_PASSWORD', ''),
        os.getenv('PINENTRY_FALLBACK_APP'),
        unknown_args,
    )
