#!/usr/bin/env python3
# Load ZFS encryption keys from remote HTTPS sources defined in /etc/zfs/remote-keys/*.json.
# Each JSON file has the format: {"dataset":"pool/name","url":"https://...","insecure":false}
# This script is invoked by zfs-remote-unlock.service at boot after network and pool import.

import glob
import json
import logging
import os
import ssl
import subprocess
import sys
import time
import urllib.error
import urllib.request

META_DIR = '/etc/zfs/remote-keys'
LOG_FILE = '/var/log/zfs-remote-unlock.log'

# Log to both stdout (picked up by journald) and a persistent logfile.
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s zfs-remote-unlock: %(levelname)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler(LOG_FILE),
    ],
)
log = logging.getLogger()


def load_config(path):
    """Parse a remote-key JSON config file. Returns (dataset, url, insecure, decrypt_passphrase, network_timeout) or raises."""
    with open(path) as f:
        data = json.load(f)
    dataset = data['dataset']
    url = data['url']
    insecure = bool(data.get('insecure', False))
    decrypt_passphrase = data.get('decrypt_passphrase', '') or ''
    network_timeout = int(data.get('network_timeout', 60))
    network_timeout = max(10, min(600, network_timeout))
    return dataset, url, insecure, decrypt_passphrase, network_timeout


def decrypt_key(key_bytes, passphrase):
    """Decrypt an openssl-aes-256-cbc-pbkdf2-base64 blob using the given passphrase."""
    result = subprocess.run(
        ['openssl', 'enc', '-d', '-aes-256-cbc', '-pbkdf2', '-base64',
         '-pass', 'pass:' + passphrase],
        input=key_bytes,
        capture_output=True,
    )
    if result.returncode != 0:
        stderr = result.stderr.decode(errors='replace').strip()
        raise RuntimeError(f'openssl decryption failed: {stderr}')
    return result.stdout


NETWORK_RETRY_DELAY = 3  # seconds


def fetch_key(url, insecure, network_timeout=60):
    """Fetch the key from the given URL, retrying on network errors. Returns the raw key bytes."""
    ctx = ssl.create_default_context()
    if insecure:
        log.warning('TLS certificate verification disabled for %s', url)
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
    req = urllib.request.Request(url, headers={'User-Agent': 'zfs-remote-unlock/1.0'})
    retries = max(1, network_timeout // NETWORK_RETRY_DELAY)
    last_exc = None
    for attempt in range(1, retries + 1):
        try:
            with urllib.request.urlopen(req, context=ctx, timeout=30) as resp:
                return resp.read()
        except urllib.error.HTTPError:
            raise  # HTTP errors (4xx/5xx) are not network problems, don't retry
        except Exception as e:
            last_exc = e
            if attempt < retries:
                log.warning('Network not ready (attempt %d/%d): %s, retrying in %ds...',
                            attempt, retries, e, NETWORK_RETRY_DELAY)
                time.sleep(NETWORK_RETRY_DELAY)
    raise last_exc


def zfs_load_key(dataset, key_bytes):
    """Pipe key_bytes into `zfs load-key <dataset>`. Returns True on success."""
    result = subprocess.run(
        ['zfs', 'load-key', dataset],
        input=key_bytes,
        capture_output=True,
    )
    if result.returncode != 0:
        stderr = result.stderr.decode(errors='replace').strip()
        log.error('zfs load-key %s failed (rc=%d): %s', dataset, result.returncode, stderr)
        return False
    return True


def zfs_mount(dataset):
    """Mount the dataset, ignoring errors (e.g. already mounted)."""
    result = subprocess.run(
        ['zfs', 'mount', dataset],
        capture_output=True,
    )
    if result.returncode != 0:
        stderr = result.stderr.decode(errors='replace').strip()
        log.warning('zfs mount %s: %s', dataset, stderr)


def main():
    if not os.path.isdir(META_DIR):
        log.debug('No remote-keys directory at %s, exiting.', META_DIR)
        sys.exit(0)

    json_files = sorted(glob.glob(os.path.join(META_DIR, '*.json')))
    if not json_files:
        log.debug('No *.json files in %s, exiting.', META_DIR)
        sys.exit(0)

    log.info('Found %d remote-key config(s) to process.', len(json_files))
    errors = 0

    for path in json_files:
        log.debug('Reading config: %s', path)
        try:
            dataset, url, insecure, decrypt_passphrase, network_timeout = load_config(path)
        except Exception as e:
            log.error('Failed to parse %s: %s', path, e)
            errors += 1
            continue

        log.info('Fetching key for %s from %s (network timeout %ds)', dataset, url, network_timeout)
        try:
            key_bytes = fetch_key(url, insecure, network_timeout)
        except urllib.error.HTTPError as e:
            log.error('HTTP error fetching key for %s: %s %s', dataset, e.code, e.reason)
            errors += 1
            continue
        except Exception as e:
            log.error('Failed to fetch key for %s: %s', dataset, e)
            errors += 1
            continue

        if decrypt_passphrase:
            log.debug('Decrypting key for %s', dataset)
            try:
                key_bytes = decrypt_key(key_bytes, decrypt_passphrase)
            except Exception as e:
                key_bytes = b''
                log.error('Failed to decrypt key for %s: %s', dataset, e)
                errors += 1
                continue

        log.debug('Key received (%d bytes), loading into %s', len(key_bytes), dataset)
        if zfs_load_key(dataset, key_bytes):
            key_bytes = b''  # clear from memory immediately after use
            log.info('Key loaded for %s, mounting...', dataset)
            zfs_mount(dataset)
            log.info('Unlocked and mounted %s', dataset)
        else:
            key_bytes = b''  # clear on failure too
            log.error('Failed to load key for %s (wrong key or already unlocked?)', dataset)
            errors += 1

    if errors:
        log.warning('%d error(s) occurred during remote unlock.', errors)
        sys.exit(1)
    else:
        log.info('All datasets unlocked successfully.')


if __name__ == '__main__':
    main()
