import zipfile
import struct
import sys
import os

def check_so_alignment(aab_path):
    print(f"Checking AAB: {aab_path}")
    
    try:
        with zipfile.ZipFile(aab_path, 'r') as z:
            so_files = [f for f in z.namelist() if f.endswith('.so') and 'arm64-v8a' in f]
            
            if not so_files:
                print("No arm64-v8a .so files found in the AAB.")
                return

            print(f"Found {len(so_files)} arm64-v8a .so files.")
            
            issues = []
            
            for so_name in so_files:
                with z.open(so_name) as f:
                    data = f.read()
                    
                    # Check ELF magic
                    if data[:4] != b'\x7fELF':
                        print(f"Skipping {so_name}: Not an ELF file")
                        continue
                    
                    # Check class (64-bit)
                    if data[4] != 2:
                        print(f"Skipping {so_name}: Not a 64-bit ELF")
                        continue

                    # Parse ELF Header (Elf64_Ehdr)
                    # e_phoff is at offset 32 (0x20), 64-bit unsigned
                    e_phoff = struct.unpack('<Q', data[32:40])[0]
                    # e_phentsize is at offset 54 (0x36), 16-bit unsigned
                    e_phentsize = struct.unpack('<H', data[54:56])[0]
                    # e_phnum is at offset 56 (0x38), 16-bit unsigned
                    e_phnum = struct.unpack('<H', data[56:58])[0]

                    max_align = 0
                    is_aligned = True
                    
                    for i in range(e_phnum):
                        offset = e_phoff + i * e_phentsize
                        # Parse Program Header (Elf64_Phdr)
                        # p_type is at offset 0, 32-bit unsigned
                        p_type = struct.unpack('<I', data[offset:offset+4])[0]
                        
                        # PT_LOAD is 1
                        if p_type == 1:
                            # p_align is at offset 48 (0x30), 64-bit unsigned
                            p_align = struct.unpack('<Q', data[offset+48:offset+56])[0]
                            max_align = max(max_align, p_align)
                            
                            if p_align < 0x4000:
                                is_aligned = False
                    
                    if not is_aligned:
                        issues.append((so_name, max_align))
                        print(f"[FAIL] {so_name}: Max alignment is {max_align} (needs 16384/0x4000)")
                    else:
                        # print(f"[PASS] {so_name}: Aligned to {max_align}")
                        pass

            if issues:
                print("\nSummary of non-compliant libraries (arm64-v8a):")
                for name, align in issues:
                    print(f" - {name} (Align: {align})")
                sys.exit(1)
            else:
                print("\nAll arm64-v8a shared libraries are 16KB aligned!")
                sys.exit(0)

    except Exception as e:
        print(f"Error checking AAB: {e}")
        sys.exit(1)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python check_16kb.py <path_to_aab>")
        sys.exit(1)
    
    check_so_alignment(sys.argv[1])