import pyshark
import struct
import zlib
import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

# --- Constants ---
SESSION_KEY = bytes.fromhex("fb2ed4f1837a722b149df99a4021177e")
SESSION_ID_LE = struct.pack('<I', 0x392DECAA)

os.makedirs("smm2_streams", exist_ok=True)

print("[*] Loading PCAP via PyShark...")
cap = pyshark.FileCapture('./decrypted_capture.pcapng', display_filter='udp')

success_count = 0

for i, pkt in enumerate(cap):
    try:
        udp_payload = bytes.fromhex(pkt.udp.payload.replace(':', ''))
        
        if not udp_payload.startswith(b'\x32\xab\x98\x64'):
            continue
            
        src_mac_str = pkt.wlan.sa.replace(':', '')
        src_mac = bytes.fromhex(src_mac_str)

        conn_id = udp_payload[5:6]
        header_nonce = udp_payload[8:16]
        tag = udp_payload[16:32]
        ciphertext = udp_payload[32:]

        crc = zlib.crc32(SESSION_ID_LE + src_mac)
        
        decrypted = False
        plaintext = b""
        
        for crc_endian in ['>I', '<I']:
            crc_bytes = struct.pack(crc_endian, crc)[:3]
            iv = crc_bytes + conn_id + header_nonce
            
            try:
                cipher_gcm = Cipher(algorithms.AES(SESSION_KEY), modes.GCM(iv, tag))
                decryptor = cipher_gcm.decryptor()
                plaintext = decryptor.update(ciphertext) + decryptor.finalize()
                decrypted = True
                break
            except Exception:
                continue
                
        if decrypted:
            # --- PIA Message Parser ---
            offset = 0
            msg_count = 0
            
            # State variables (inherit from previous message if missing)
            msg_flags = 0
            payload_size = 0
            protocol_type = 0
            
            while offset < len(plaintext):
                msg_count += 1
                presence_flags = plaintext[offset]
                offset += 1
                
                if presence_flags & 1:
                    msg_flags = plaintext[offset]
                    offset += 1
                    
                if presence_flags & 2:
                    payload_size = struct.unpack('>H', plaintext[offset:offset+2])[0]
                    offset += 2
                    
                if presence_flags & 4:
                    protocol_type = plaintext[offset]
                    offset += 4 # Skip protocol type (1) and port (3)
                    
                if presence_flags & 8:
                    offset += 8 # Skip destination (8)
                    
                if presence_flags & 16:
                    offset += 8 # Skip source constant ID (8)
                    
                is_compressed = bool(msg_flags & 0x10)
                
                # Extract the pure application payload
                raw_payload = plaintext[offset:offset+payload_size]
                
                if protocol_type != 0x81:
                    continue

                
                if payload_size >= 0xB:
                    if is_compressed:
                        continue
                    else:
                        with open(f"smm2_streams/pkt{i}_msg{msg_count}_proto{protocol_type:02x}.bin", "wb") as f:
                            f.write(raw_payload)
                        success_count += 1
                        
                # PIA Messages are always padded to 4-byte boundaries
                offset += payload_size
                align = (4 - (offset % 4)) % 4
                offset += align

    except AttributeError:
        continue
    except Exception as e:
        continue

cap.close()
print(f"\n[+] Finished! Extracted {success_count} payloads to 'smm2_streams' folder.")