#!/usr/bin/env python
import argparse
import binascii
import os
import platform
import shutil
import sys

from Crypto.Cipher import AES
from Crypto.Hash import HMAC, SHA256
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Util import Padding


from random import SystemRandom
random = SystemRandom()


def take(bytearr, start, count):
    return bytearr[start:start+count]


def write_file(path, content):
    if os.path.exists(path):
        randomsuffix = binascii.hexlify(random.randbytes(4)).decode()
        backupfile = path + f".{randomsuffix}.backup"
        print(f"making backup copy of {path} to {backupfile}", file=sys.stderr)
        shutil.copy(path, backupfile)
    print(f"writing data to {path}", file=sys.stderr)
    mode = "wb" if isinstance(content, bytes) else "w"
    with open(path, "wb") as f:
        f.write(content.encode() if isinstance(content, str) else content)


PASSW = "e6552ddf3ac5c8755a82870d91273a63eab0da1e"
KDF_ITER_CNT = 10000
KEY_BYTE_LEN = 32
SALT_BYTE_LEN = 8
IV_BYTE_LEN = 16
NOTSECRETPAYLOAD = b"4af85e84255b077ad890dba297e811b7d016add1"
HMAC_BYTE_LEN = 32


def encrypt(data_raw):
    encoded_data = NOTSECRETPAYLOAD
    cryptsalt = random.randbytes(SALT_BYTE_LEN)
    encoded_data += cryptsalt
    authsalt = random.randbytes(SALT_BYTE_LEN)
    encoded_data += authsalt

    cryptkey = PBKDF2(PASSW, cryptsalt, KEY_BYTE_LEN, KDF_ITER_CNT)
    authkey = PBKDF2(PASSW, authsalt, KEY_BYTE_LEN, KDF_ITER_CNT)

    iv = random.randbytes(IV_BYTE_LEN)
    encoded_data += iv
    ciphertext = AES.new(cryptkey, AES.MODE_CBC, iv).encrypt(
        Padding.pad(data_raw, block_size=AES.block_size)
    )
    encoded_data += ciphertext

    tag = HMAC.new(authkey, encoded_data, SHA256).digest()
    return encoded_data + tag


def decrypt(data_raw):
    notsecret_len = len(NOTSECRETPAYLOAD)
    cryptsalt = take(data_raw, notsecret_len, SALT_BYTE_LEN)
    # authsalt = take(data_raw, NONSECRET_LEN+SALT_BYTE_LEN, SALT_BYTE_LEN)
    cryptkey = PBKDF2(PASSW, cryptsalt, KEY_BYTE_LEN, KDF_ITER_CNT)
    # authkey = PBKDF2(PASSW, authsalt, KEY_BYTE_LEN, KDF_ITER_CNT)
    iv = take(data_raw, notsecret_len + SALT_BYTE_LEN * 2, IV_BYTE_LEN)
    ciphertext_start = notsecret_len + SALT_BYTE_LEN * 2 + IV_BYTE_LEN
    ciphertext = data_raw[ciphertext_start:-HMAC_BYTE_LEN]
    return Padding.unpad(
        AES.new(cryptkey, AES.MODE_CBC, iv).decrypt(ciphertext),
        block_size=AES.block_size
    )


HEADER_LEN = 3
PROFILE_ID_LEN = 64
PLAINTEXT_PROFILE_TYPE = b"v2n"
XML_HEADER_START = b'<?xml'


def encode(profile_id, data_raw):
    assert data_raw.startswith(XML_HEADER_START), \
        f"Encoding failed, expected XML header {XML_HEADER_START}, " \
        f"got {data_raw[:len(XML_HEADER_START)]}"
    encoded = PLAINTEXT_PROFILE_TYPE
    encoded += profile_id.encode()
    encoded += encrypt(data_raw)

    return encoded


def decode(data_raw):
    profile_type = take(data_raw, 0, HEADER_LEN)
    assert profile_type == PLAINTEXT_PROFILE_TYPE, \
        f"Profile type should be {PLAINTEXT_PROFILE_TYPE}, got " \
        f"{profile_type}"

    profile_id = take(data_raw, HEADER_LEN, PROFILE_ID_LEN)
    print(f"profile id: {profile_id}", file=sys.stderr)
    plaintext = decrypt(
        data_raw[HEADER_LEN+PROFILE_ID_LEN:]
    )
    assert plaintext.startswith(XML_HEADER_START), \
        f"Decoding failed, expected XML header {XML_HEADER_START}, " \
        f"got {plaintext[:len(XML_HEADER_START)]}"
    return plaintext.decode()


def decrypt_main(args):
    with open(args.infile, "rb") as f:
        data_raw = f.read()
    decoded = decode(data_raw)
    print("decoding successful", file=sys.stderr)
    if args.outfile:
        write_file(args.outfile, decoded)
    else:
        print(f"profile content ({len(decoded)} characters):\n", file=sys.stderr)
        print(decoded)


def fetch_profile_id(path):
    assert os.path.exists(path), \
        "profile id was not provided, and " \
        f"also cannot be read from {path} (because it doesn't exist)"

    print(f"trying to fetch profile id from {path}", file=sys.stderr)
    with open(path, "rb") as f:
        f.read(HEADER_LEN)
        data_raw = f.read(PROFILE_ID_LEN)
    assert len(data_raw) == PROFILE_ID_LEN, \
        f"Couldn't read {PROFILE_ID_LEN} character profile id from {path} " \
        "- malformed file?"
    return data_raw.decode()


def encrypt_main(args):
    profile_id = args.profile_id or fetch_profile_id(args.outfile)
    try:
        binascii.unhexlify(profile_id)
    except Exception as ex:
       raise ValueError(f"profile id ({profile_id}) is invalid") from ex
    with open(args.infile, "rb") as f:
        data_raw = f.read()
    encoded = encode(profile_id, data_raw)
    write_file(args.outfile, encoded)


DEFAULT_PROFILE_PATH = os.path.expanduser(
    "~\\AppData\\Local\\Eddie\\"
    if platform.system().lower() == "windows"
    else '~/.config/eddie/'
) + "default.profile"


def parse_args():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(required=True, help='encrypt or decrypt')
    e_parser = subparsers.add_parser("encrypt", aliases=["-e", "e"])
    e_parser.add_argument(
        "-i", "--infile",
        required=True,
        help='File path to read plain XML text data from'
    )
    e_parser.add_argument(
        "-o", "--outfile",
        default=DEFAULT_PROFILE_PATH,
        help='File path to write the encoded profile to'
    )
    e_parser.add_argument(
        "-p", "--profile-id",
        help='64 hexadecimal character string, by default it will be fetched from OUTFILE'
    )
    e_parser.set_defaults(action=encrypt_main)

    d_parser = subparsers.add_parser("decrypt", aliases=["-d", "d"])
    d_parser.add_argument(
        "-i", "--infile", default=DEFAULT_PROFILE_PATH,
        help="File path to read profile from"
    )
    d_parser.add_argument(
        "-o", "--outfile",
        help="File path to write to, by default it's stdout"
    )
    d_parser.set_defaults(action=decrypt_main)

    return parser.parse_args()


def main():
    args = parse_args()
    args.action(args)


if __name__ == '__main__':
    main()