#!/usr/bin/env python3
"""Decrypt CloudChute files with encryption_config.json. Does not print the master key."""

import argparse
import base64
import getpass
import hashlib
import json
import sys
from pathlib import Path

from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

MAGIC = bytes([0x43, 0x4C, 0x44, 0x43, 0x48, 0x01])
NONCE_LEN = 12
TAG_LEN = 16


def fail(message):
    print(message, file=sys.stderr)
    raise SystemExit(1)


def load_config(path):
    try:
        config = json.loads(Path(path).read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as error:
        fail(f"Could not read config: {error}")
    if config.get("version") != 1:
        fail(f"Unsupported config version: {config.get('version')!r}")
    kdf = config.get("kdfParams") or {}
    if kdf.get("algorithm") != "PBKDF2-SHA512":
        fail(f"Unsupported KDF: {kdf.get('algorithm')!r}")
    if kdf.get("keySize") != 256:
        fail(f"Unsupported key size: {kdf.get('keySize')!r}")
    if not isinstance(kdf.get("rounds"), int) or kdf["rounds"] < 1:
        fail("Config is missing kdfParams.rounds")
    for field in ("salt", "encryptedMasterKey", "recoveryMasterKey", "checkHash"):
        if not isinstance(config.get(field), str):
            fail(f"Config is missing {field}")
    return config


def b64(value, label):
    try:
        return base64.b64decode(value, validate=True)
    except Exception:
        fail(f"Config field {label} is not valid Base64")


def open_gcm(key, combined):
    if len(combined) < NONCE_LEN + TAG_LEN:
        raise ValueError("AES-GCM box is too short")
    nonce, body = combined[:NONCE_LEN], combined[NONCE_LEN:]
    try:
        return AESGCM(key).decrypt(nonce, body, None)
    except InvalidTag:
        raise ValueError("AES-GCM authentication failed")


def master_key(config, secret, field):
    salt = b64(config["salt"], "salt")
    print(f"Deriving key ({config['kdfParams']['rounds']} rounds)...", file=sys.stderr)
    wrapping_key = PBKDF2HMAC(
        algorithm=hashes.SHA512(),
        length=32,
        salt=salt,
        iterations=config["kdfParams"]["rounds"],
    ).derive(secret.encode("utf-8"))
    wrapped = b64(config[field], field)
    try:
        mek = open_gcm(wrapping_key, wrapped)
    except ValueError:
        fail("Could not unwrap the master key. Check the password or recovery phrase.")
    expected = b64(config["checkHash"], "checkHash")
    if hashlib.sha256(mek).digest() != expected:
        fail("Unwrapped key does not match checkHash.")
    return mek


def plaintext_bytes(data, mek):
    if not data.startswith(MAGIC):
        return data, False
    return open_gcm(mek, data[len(MAGIC):]), True


def write_bytes(destination, data, force):
    if destination.exists() and not force:
        raise FileExistsError(f"refusing to overwrite {destination} (pass --force)")
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_bytes(data)


def decrypt_file(source, destination, mek, force):
    data = source.read_bytes()
    plain, encrypted = plaintext_bytes(data, mek)
    write_bytes(destination, plain, force)
    return encrypted


def iter_files(root, output_root):
    for path in sorted(root.rglob("*")):
        if not path.is_file():
            continue
        if output_root is not None:
            try:
                path.resolve().relative_to(output_root.resolve())
                continue
            except ValueError:
                pass
        yield path


def prompt_secret(use_recovery):
    if use_recovery:
        raw = getpass.getpass("Recovery phrase: ")
        words = raw.split()
        if len(words) < 12:
            fail("Phrase is too short (expected 24 words).")
        return " ".join(words), "recoveryMasterKey"
    password = getpass.getpass("Master password: ")
    if password == "":
        fail("Password was empty.")
    return password, "encryptedMasterKey"


def main():
    parser = argparse.ArgumentParser(description="Decrypt CloudChute backup files.")
    parser.add_argument("config", help="Path to encryption_config.json")
    parser.add_argument("path", help="A downloaded file or folder")
    parser.add_argument("--output", help="Destination file (for a file) or folder (for a folder)")
    parser.add_argument("--recovery", action="store_true", help="Use the 24-word recovery phrase")
    parser.add_argument("--force", action="store_true", help="Overwrite an existing destination")
    args = parser.parse_args()

    config = load_config(args.config)
    source = Path(args.path)
    if not source.exists():
        fail(f"Not found: {source}")
    secret, field = prompt_secret(args.recovery)
    mek = master_key(config, secret, field)

    failures = 0
    if source.is_file():
        destination = Path(args.output) if args.output else source.with_name(source.name + ".decrypted")
        try:
            encrypted = decrypt_file(source, destination, mek, args.force)
        except (OSError, ValueError) as error:
            fail(f"{source}: {error}")
        print(("Decrypted" if encrypted else "Copied") + f": {destination}", flush=True)
        return

    if not source.is_dir():
        fail(f"Not a file or folder: {source}")
    output_root = Path(args.output) if args.output else source.with_name(source.name + "-decrypted")
    for path in iter_files(source, output_root):
        relative = path.relative_to(source)
        destination = output_root / relative
        try:
            encrypted = decrypt_file(path, destination, mek, args.force)
        except (OSError, ValueError) as error:
            failures += 1
            print(f"Failed: {relative}: {error}", file=sys.stderr, flush=True)
            continue
        print(("Decrypted" if encrypted else "Copied") + f": {relative}", flush=True)
    if failures:
        fail(f"{failures} file(s) failed.")


if __name__ == "__main__":
    main()