kuma_packages_tool.py

#!/usr/bin/env python3
# pip install pycryptodome pymongo

import argparse
import gzip
import json
import os
import sys

import bson
from Crypto.Cipher import AES
from Crypto.Hash import SHA256


def aes_encrypt(data: bytes, key: bytes) -> bytes:
    cipher = AES.new(key, AES.MODE_GCM, nonce=os.urandom(12))
    ciphertext, tag = cipher.encrypt_and_digest(data)
    return cipher.nonce + ciphertext + tag


def aes_decrypt(data: bytes, key: bytes) -> bytes:
    nonce, authtag = data[:12], data[-16:]
    cipher = AES.new(key, AES.MODE_GCM, nonce)
    return cipher.decrypt_and_verify(data[12:-16], authtag)


def decode_bson(data: bytes) -> dict:
    # Newer KUMA versions (4.2+) compress BSON with gzip before encryption
    if data[:2] == b"\x1f\x8b":
        data = gzip.decompress(data)
    return bson.decode(data)


def encode_bson(data: dict, compress: bool = False) -> bytes:
    encoded = bson.encode(data)
    if compress:
        encoded = gzip.compress(encoded)
    return encoded


def make_key(password: str) -> bytes:
    h = SHA256.new()
    h.update(password.encode())
    return h.digest()


def decrypt(input_file: str, output_file: str, key: bytes, pretty: bool) -> None:
    with open(input_file, "rb") as f:
        ciphertext = f.read()
    try:
        decrypted_data = aes_decrypt(ciphertext, key)
    except ValueError:
        print("Error while decrypt: MAC check failed\nCheck your password", file=sys.stderr)
        sys.exit(1)
    try:
        json_data = decode_bson(decrypted_data)
    except Exception as e:
        print(f"Error while decoding BSON: {e}", file=sys.stderr)
        sys.exit(1)

    class BytesEncoder(json.JSONEncoder):
        def default(self, obj):
            if isinstance(obj, (bytes, bytearray)):
                return obj.decode("utf-8", errors="replace")
            return super().default(obj)

    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(json_data, f, cls=BytesEncoder, indent=2 if pretty else None, ensure_ascii=False)


def encrypt(input_file: str, output_file: str, key: bytes, compress: bool) -> None:
    with open(input_file, "r", encoding="utf-8") as f:
        json_data = json.load(f)
    try:
        bson_data = encode_bson(json_data, compress=compress)
    except Exception as e:
        print(f"Error while encoding BSON: {e}", file=sys.stderr)
        sys.exit(1)
    encrypted_data = aes_encrypt(bson_data, key)
    with open(output_file, "wb") as f:
        f.write(encrypted_data)


def main():
    parser = argparse.ArgumentParser(
        description="Encrypt/decrypt KUMA packages",
        usage="python3 kuma_package.py [-h] [-d | -e] [-z] -p PASSWORD -f FILE -o FILE [--pretty]",
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument("-d", action="store_true", help="decrypt package")
    group.add_argument("-e", action="store_true", help="encrypt package")
    parser.add_argument("-p", metavar="PASSWORD", required=True, help="password")
    parser.add_argument("-f", metavar="FILE", required=True, help="input file")
    parser.add_argument("-o", metavar="FILE", required=True, help="output file")
    parser.add_argument("--pretty", action="store_true", help="human readable format with indents")
    parser.add_argument(
        "-z",
        action="store_true",
        help="gzip-compress before encrypting (required for KUMA 4.1+, auto-detected on decrypt)",
    )
    args = parser.parse_args()

    if not args.d and not args.e:
        parser.error("specify -d (decrypt) or -e (encrypt)")

    key = make_key(args.p)

    if args.d:
        decrypt(args.f, args.o, key, args.pretty)
    else:
        encrypt(args.f, args.o, key, compress=args.z)


if __name__ == "__main__":
    main()