#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import io
import struct
import sys
import os

ZIP_EOCD_SIGNATURE = b'\x50\x4b\x05\x06'  # End of Central Directory 시그니처 (0x06054B50)
CENTRAL_DIRECTORY_SIGNATURE = b'\x50\x4b\x01\x02'  # Central Directory 시그니처 (0x02014B50)
CENTRAL_DIRECTORY_BASE_SIZE = 46  # FileName 제외한 Central Directory 구조체 크기 (고정)

def find_eocd_offset(data: bytes) -> int:
    max_comment_size = 65535
    max_eocd_search = max_comment_size + 22

    filesize = len(data)
    search_size = min(filesize, max_eocd_search)
    # 파일 끝에서부터 search_size만큼 잘라내어 End of Central Directory 시그니처를 rfind
    window = data[filesize - search_size:]
    pos = window.rfind(ZIP_EOCD_SIGNATURE)
    if pos == -1:
        raise ValueError("End of Central Directory 시그니처(0x06054B50)를 찾을 수 없습니다.")
    return (filesize - search_size) + pos

def patch_entry_encrypted(input_file: str) -> list:
    # 파일 전체를 메모리에 로드
    with open(input_file, "rb") as f:
        data = f.read()

    # End of Central Directory 오프셋 찾기
    eocd_offset = find_eocd_offset(data)

    with io.BytesIO(bytearray(data)) as zf:
        # End of Central Directory 시그니처 확인
        zf.seek(eocd_offset, 0)
        eocd_signature = zf.read(4)
        if eocd_signature != ZIP_EOCD_SIGNATURE:
            raise ValueError("End of Central Directory 시그니처가 올바르지 않습니다.")

        # End of Central Directory 구조체 파싱 (최소 22바이트)
        zf.seek(eocd_offset, 0)
        eocd_data = zf.read(22)
        eocd_struct = struct.unpack("<IHHHHLLH", eocd_data)  
        
        # 필요한 End of Central Directory 필드만 변수에 매핑
        _, _, _, disk_entries, total_entries, central_directory_size, central_directory_offset, comment_len = eocd_struct

        # 중앙 디렉터리(Central Directory)로 이동
        zf.seek(central_directory_offset, 0)

        patched_info = []
        for _ in range(total_entries):
            cd_offset = zf.tell()
            cd_data = zf.read(CENTRAL_DIRECTORY_BASE_SIZE)
            
            # Central Directory 구조체(46바이트 고정) 파싱
            cd_struct = struct.unpack("<IHHHHHHIIIHHHHHII", cd_data)
            signature_cd        = cd_struct[0]
            version_by          = cd_struct[1]
            version_needed      = cd_struct[2]
            flags               = cd_struct[3]
            compression         = cd_struct[4]
            time                = cd_struct[5]
            date                = cd_struct[6]
            crc32               = cd_struct[7]
            comp_size           = cd_struct[8]
            uncomp_size         = cd_struct[9]
            filename_length     = cd_struct[10]
            extra_length        = cd_struct[11]
            comment_length      = cd_struct[12]
            disk_num_start      = cd_struct[13]
            internal_attr       = cd_struct[14]
            external_attr       = cd_struct[15]
            local_header_offset = cd_struct[16]

            # 파일명 추출
            filename = zf.read(filename_length).decode(errors='replace')
            # Extra, Comment는 길이만큼 Skip
            zf.seek(extra_length + comment_length, 1)

            # 암호화 비트가 켜져 있다면 해제
            if (flags & 0x0001) != 0:
                old_flags = flags
                new_flags_value = flags & ~0x0001  # 비트 0 제거
                new_flags = struct.pack("<H", new_flags_value)

                # CD 상에서 flags 오프셋은 cd_offset + 8 바이트 지점
                flags_offset = cd_offset + 8

                # 현재 지점을 저장 후 flags 오프셋에 new_flags 기록
                cur_pos = zf.tell()
                zf.seek(flags_offset)
                zf.write(new_flags)
                zf.seek(cur_pos)

                patched_info.append({
                    "FileName": filename,
                    "PatchOffset": flags_offset,
                    "OldFlags": old_flags,
                    "NewFlags": new_flags_value
                })

        # 패치된 파일 쓰기
        output_file = os.path.splitext(input_file)[0] + "_patched.apk"
        with open(output_file, "wb") as out:
            out.write(zf.getvalue())

    return patched_info

def main():
    if len(sys.argv) < 2:
        print(f"Usage: {os.path.basename(sys.argv[0])} <target_apk_file>")
        sys.exit(1)

    input_file = sys.argv[1]
    if not os.path.exists(input_file):
        print(f"[!] 파일이 존재하지 않습니다: {input_file}")
        sys.exit(1)

    try:
        patched_files = patch_entry_encrypted(input_file)
    except Exception as e:
        print(f"[!] 패치 중 오류가 발생했습니다: {e}")
        sys.exit(1)

    # 결과 로그 출력 및 파일 기록
    log_file = input_file + ".patch.log"
    with open(log_file, "w", encoding="utf-8") as lf:
        for info in patched_files:
            log_line = (f"FileName: {info['FileName']}, "
                        f"PatchOffset: 0x{info['PatchOffset']:08X}, "
                        f"OldFlags: 0x{info['OldFlags']:04X}, "
                        f"NewFlags: 0x{info['NewFlags']:04X}")
            print(log_line)
            lf.write(log_line + "\n")

    print(f"\n[+] 총 {len(patched_files)}개 파일의 암호화 플래그를 해제했습니다.")
    print(f"[+] 패치된 APK: {os.path.splitext(input_file)[0]}_patched.apk")
    print(f"[+] 패치 로그 파일: {log_file}")

if __name__ == "__main__":
    main()