import hashlib
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Cipher import AES
from Crypto.Hash import SHA1
from .as_req import build_as_req, build_full_as_req
from .kdc import send_kerberos_packet
from .asn1 import get_asn1_len
from .tools import parse_krb_error


def extract_kerberos_salt(error_data: bytes) -> str:
    if error_data[0] != 0x7e:
        raise ValueError("Invalid KRB-ERROR packet")

    # Tag 12 (e-data) is [12] EXPLICIT OCTET STRING or SEQUENCE
    e_data_idx = error_data.find(b'\xac')
    if e_data_idx == -1:
        raise ValueError("e-data (Tag 12) not found")

    # The Salt is a GeneralString (Tag 0x1b)
    salt_tag_idx = error_data.find(b'\x1b', e_data_idx)
    if salt_tag_idx == -1:
        raise ValueError("Salt (GeneralString Tag 0x1b) not found in e-data")

    salt_len, salt_len_size = get_asn1_len(error_data, salt_tag_idx + 1)
    salt_start = salt_tag_idx + 1 + salt_len_size

    salt_bytes = error_data[salt_start: salt_start + salt_len]
    return salt_bytes.decode('utf-8')


def nfold_for_kerberos():
    return b'\x6b\x65\x72\x62\x65\x72\x6f\x73\x7b\x9b\x5b\x2b\x93\x13\x2b\x93'


def aes_string_to_key(password, salt):
    tkey = PBKDF2(
        password.encode(),
        salt.encode(),
        dkLen=32,
        count=4096,
        hmac_hash_module=SHA1
    )
    constant = nfold_for_kerberos()

    iv = b'\x00' * 16
    cipher1 = AES.new(tkey, AES.MODE_CBC, iv)
    part1 = cipher1.encrypt(constant)

    cipher2 = AES.new(tkey, AES.MODE_CBC, iv)
    part2 = cipher2.encrypt(part1)

    return part1 + part2


async def get_tgt(username: str, password: str, domain: str,
                  kdc_host: str, kdc_port: int = 88) -> tuple[bytes, bytes]:
    as_req = build_as_req(username, domain)
    resp = await send_kerberos_packet(as_req, kdc_host, kdc_port)
    parse_krb_error(resp)
    salt = extract_kerberos_salt(resp)
    base_key = aes_string_to_key(password, salt)
    full_as_req = build_full_as_req(username, domain, base_key)
    as_res_bytes = await send_kerberos_packet(full_as_req, kdc_host, kdc_port)
    parse_krb_error(as_res_bytes)
    return as_res_bytes, base_key
