# Copy this file to ~/.ansible/plugins/lookup/keepass.py
#
# Prerequisites:
#   pip install keepassxc-proxy-client
#
# KeePassXC setup:
#   Tools > Settings > Browser Integration > Enable browser integration
#   The plugin will prompt you to approve the association in KeePassXC on first run.
#   Association credentials are saved to ~/.keepassxc_ansible_assoc (override with
#   env var KEEPASSXC_ASSOC_FILE).
#
# Usage:
#   {{ lookup('keepass', 'https://example.com', 'password') }}
#   {{ lookup('keepass', 'https://example.com', 'username') }}
#   {{ lookup('keepass', 'https://example.com') }}          # defaults to password
#
# Terms:
#   terms[0]  URL (or title) to search for — matched against KeePassXC entries
#   terms[1]  Attribute to return: username / user / password / pass / passwd
#             Defaults to 'password' if omitted.

from __future__ import absolute_import, division, print_function
__metaclass__ = type

import json
import os

from ansible.errors import AnsibleError
from ansible.plugins.lookup import LookupBase

DOCUMENTATION = r"""
    lookup: keepass
    short_description: Fetch credentials from KeePassXC via the browser proxy protocol
    description:
        - Connects to a running KeePassXC instance over its browser-integration Unix socket
          (or Windows named pipe) and retrieves credentials for a given URL.
        - On first use the plugin associates itself with KeePassXC; you will see an
          approval dialog in KeePassXC. The association is saved to disk so subsequent
          runs are silent.
    options:
      _terms:
        description:
          - First term is the URL (or entry title URL) to look up.
          - Second term (optional) is the attribute to return.
            Accepted values are username, user, password, pass, passwd.
            Defaults to password.
        required: true
    notes:
      - KeePassXC must be running with Browser Integration enabled.
      - Requires the Python package keepassxc-proxy-client.
"""

EXAMPLES = r"""
- name: Get password for an entry
  debug:
    msg: "{{ lookup('keepass', 'https://example.com') }}"

- name: Get username for an entry
  debug:
    msg: "{{ lookup('keepass', 'https://example.com', 'username') }}"

- name: Set pi account password from KeePassXC
  user:
    name: pi
    password: "{{ lookup('keepass', 'default_rpi_password') | password_hash('sha512') }}"
"""

ASSOC_FILE_ENV = "KEEPASSXC_ASSOC_FILE"
ASSOC_FILE_DEFAULT = os.path.expanduser("~/.keepassxc_ansible_assoc")
ASSOC_NAME = "ansible-keepassxc"

USERNAME_ATTRS = {"username", "user"}
PASSWORD_ATTRS = {"password", "pass", "passwd"}


def _get_assoc_file():
    return os.environ.get(ASSOC_FILE_ENV, ASSOC_FILE_DEFAULT)


def _load_assoc(connection):
    """Load a previously saved association from disk into *connection*. Returns True on success."""
    path = _get_assoc_file()
    if not os.path.exists(path):
        return False
    try:
        with open(path, "r") as fh:
            data = json.load(fh)
        connection.load_associate(data["name"], bytes(data["public_key"]))
        return True
    except (KeyError, ValueError, OSError):
        return False


def _save_assoc(connection):
    """Persist the current association to disk."""
    path = _get_assoc_file()
    assoc_id, public_key = connection.dump_associate()
    data = {"name": assoc_id, "public_key": list(public_key)}
    with open(path, "w") as fh:
        json.dump(data, fh)
    os.chmod(path, 0o600)


def _connect():
    """Return an authenticated, test-verified Connection to KeePassXC."""
    try:
        from keepassxc_proxy_client.protocol import Connection, ResponseUnsuccesfulException
    except ImportError:
        raise AnsibleError(
            "keepassxc-proxy-client is not installed. "
            "Run: pip install keepassxc-proxy-client"
        )

    conn = Connection()

    try:
        conn.connect()
    except Exception as exc:
        raise AnsibleError(
            "Cannot connect to KeePassXC. Is it running with Browser Integration enabled? "
            "Error: {}".format(exc)
        )

    conn.change_public_keys()

    # Try to reuse an existing association
    if _load_assoc(conn):
        try:
            conn.test_associate(trigger_unlock=True)
            return conn  # existing association is still valid
        except Exception:
            pass  # fall through and re-associate

    # First run (or stale association) — ask KeePassXC to authorise us
    try:
        conn.associate()
    except Exception as exc:
        raise AnsibleError(
            "KeePassXC association failed. "
            "Please approve the request in the KeePassXC dialog. "
            "Error: {}".format(exc)
        )

    try:
        conn.test_associate(trigger_unlock=True)
    except Exception as exc:
        raise AnsibleError("KeePassXC association test failed: {}".format(exc))

    _save_assoc(conn)
    return conn


class LookupModule(LookupBase):

    def run(self, terms, variables=None, **kwargs):
        if not terms or len(terms) > 2:
            raise AnsibleError(
                "keepass lookup expects 1 or 2 arguments: "
                "lookup('keepass', '<url>', '<attribute>')"
            )

        url = terms[0]
        attr = terms[1].lower() if len(terms) == 2 else "password"

        if attr not in USERNAME_ATTRS | PASSWORD_ATTRS:
            raise AnsibleError(
                "keepass: unsupported attribute '{}'. "
                "Use one of: username, user, password, pass, passwd".format(attr)
            )

        conn = _connect()

        try:
            entries = conn.get_logins(url)
        except Exception as exc:
            raise AnsibleError(
                "keepass: failed to retrieve logins for '{}': {}".format(url, exc)
            )

        if not entries:
            raise AnsibleError(
                "keepass: no entries found in KeePassXC for URL '{}'".format(url)
            )

        # Return the first matching entry (KeePassXC already filters by URL)
        entry = entries[0]

        if attr in USERNAME_ATTRS:
            return [entry.get("login", "")]
        else:
            return [entry.get("password", "")]