#!/usr/bin/python
# -*- coding: utf-8 -*-

import argparse
import base64
import codecs
import sqlite3
import json
import sys

from Cryptodome import Random
from Cryptodome.Cipher import AES
from Cryptodome.Util import Counter

def rclone_encrypt_password(password):
    key = bytes([0x9c, 0x93, 0x5b, 0x48, 0x73, 0x0a, 0x55, 0x4d,
                 0x6b, 0xfd, 0x7c, 0x63, 0xc8, 0x86, 0xa9, 0x2b,
                 0xd3, 0x90, 0x19, 0x8e, 0xb8, 0x12, 0x8a, 0xfb,
                 0xf4, 0xde, 0x16, 0x2b, 0x8b, 0x95, 0xf6, 0x38])

    nonce = Random.new().read(AES.block_size)
    counter = Counter.new(128, initial_value=int(codecs.encode(nonce, "hex"), 16))
    cipher = AES.new(key, AES.MODE_CTR, counter=counter)

    encrypted_bytes = nonce + cipher.encrypt(password.encode("utf-8"))
    return base64.urlsafe_b64encode(encrypted_bytes).decode("ascii").rstrip("=")

def decrypt_data(encrypted_b64_string, secret_key):
    if not encrypted_b64_string:
        return b''

    encrypted_bytes = base64.b64decode(encrypted_b64_string)
    nonce = encrypted_bytes[:8]
    encrypted_payload = encrypted_bytes[8:]

    cipher = AES.new(secret_key, AES.MODE_CTR, counter=Counter.new(64, prefix=nonce))

    return cipher.decrypt(encrypted_payload).rstrip(b'{').decode('utf8')

parser = argparse.ArgumentParser(
    description="Extracts credentials from TrueNAS database"
)
parser.add_argument(
    '--backup-credential-id',
    type=int,
    help='Backup credential ID'
)
parser.add_argument(
    '--id',
    type=int,
    help='Task ID for --truecloud or --cloudsync'
)
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument(
    '--truecloud',
    action='store_true',
    help='Use TrueCloud provider'
)
group.add_argument(
    '--cloudsync',
    action='store_true',
    help='Use CloudSync provider'
)
args = parser.parse_args()

if (args.truecloud or args.cloudsync) and args.id is None:
    parser.error("--id is required when using --truecloud or --cloudsync")

secret_path = '/data/pwenc_secret'
try:
    with open(secret_path, 'rb') as f:
        secret = f.read()
except FileNotFoundError:
    print(f"Error: password secret seed encryption file not found: {secret_path}", file=sys.stderr)
    exit(1)

db_path = '/data/freenas-v1.db'
backup_credentials_encrypted = None
encryption_password_encrypted = None
encryption_salt_encrypted = None

try:
    con = sqlite3.connect(db_path)
    cur = con.cursor()

    res_backup_credentials = cur.execute("SELECT attributes FROM system_cloudcredentials WHERE id = ?", (args.backup_credential_id,))
    row_backup_credentials = res_backup_credentials.fetchone()
    if row_backup_credentials:
        (backup_credentials_encrypted,) = row_backup_credentials
    else:
        print("Error: backup credentials not found", file=sys.stderr)
        sys.exit(1)

    if args.truecloud:
        res_encryption = cur.execute("SELECT password FROM tasks_cloud_backup WHERE id = ?", (args.id,))
        row_encryption = res_encryption.fetchone()
        if row_encryption:
            (encryption_password_encrypted,) = row_encryption
        else:
            print("Error: TrueCloud Restic password not found", file=sys.stderr)
            sys.exit(1)
    elif args.cloudsync:
        res_encryption = cur.execute("SELECT encryption_password, encryption_salt FROM tasks_cloudsync WHERE id = ?", (args.id,))
        row_encryption = res_encryption.fetchone()
        if row_encryption:
            (encryption_password_encrypted, encryption_salt_encrypted) = row_encryption
        else:
            print("Error: CloudSync encryption password not found", file=sys.stderr)
            sys.exit(1)

finally:
    if con:
        con.close()

decrypted_backup_credentials_json_str = decrypt_data(backup_credentials_encrypted, secret)
backup_credentials = json.loads(decrypted_backup_credentials_json_str) if decrypted_backup_credentials_json_str else {}
encryption_password = decrypt_data(encryption_password_encrypted, secret)
if args.cloudsync:
    encryption_password = rclone_encrypt_password(encryption_password)

final_data = {
    "encryption_password": encryption_password,
    "backup_credentials": backup_credentials
}

if args.cloudsync and encryption_salt_encrypted:
    encryption_salt = decrypt_data(encryption_salt_encrypted, secret)
    final_data["encryption_salt"] = rclone_encrypt_password(encryption_salt)

print(json.dumps(final_data, indent=4, ensure_ascii=False))