import argparse
import sys
import lz4.block

# NOTE: miitomo common key is 9ec1c78fa2cb34e2bed5691c08432f04

# session id is player_session_id cookie

def transform_common_key(s):
    """
    Transforms the common key by subtracting 0x62 from each character
    and negating the result, similar to FUN_0004a120 in libcocos2dcpp.so.
    
    Args:
        s (str): Input common key string.
    
    Returns:
        bytes: Processed bytes.
    """
    return bytes([(-0x62 - ord(c)) & 0xFF for c in s])

def build_xor_table(common_key, session_id):
    """
    Build the XOR table by processing the common key and session ID
    and concatenating the results.
    
    Args:
        common_key (str): The common key string.
        session_id (str): The session ID string.
    
    Returns:
        bytes: The XOR table.
    """
    transformed_common = transform_common_key(common_key)
    return transformed_common +  session_id.encode('ascii')

def conditional_xor(data, xor_table):
    """
    Apply the conditional XOR operation to the data using the XOR table.
    For each byte in the data:
        - If (key_byte & 7) == 0, XOR the data byte with the key byte.
        - Else, perform a bit rotation based on (key_byte & 7).
    
    Args:
        data (bytes): The obfuscated data.
        xor_table (bytes): The XOR table.
    
    Returns:
        bytes: The data after applying the XOR operation.
    """
    output = bytearray(len(data))
    table_len = len(xor_table)
    
    for i in range(len(data)):
        key_byte = xor_table[(i + 1) % table_len]
        
        if (key_byte & 7) == 0:
            # Perform XOR
            output[i] = data[i] ^ key_byte
        else:
            # Perform bit rotation
            shift = key_byte & 7
            rotated = ((data[i] >> (8 - shift)) | (data[i] << shift)) & 0xFF
            output[i] = rotated
    
    return bytes(output)

def decode_varint(data):
    """
    Decode a varint from the beginning of the data.
    
    Args:
        data (bytes): The data containing the varint at the start.
    
    Returns:
        tuple: (decoded integer, number of bytes consumed)
    
    Raises:
        ValueError: If the varint is too long or incomplete.
    """
    value = 0
    shift = 0
    for i, byte in enumerate(data):
        value |= (byte & 0x7F) << shift
        if (byte & 0x80) == 0:
            return value, i + 1
        shift += 7
        if shift >= 35:
            raise ValueError("Varint too long")
    raise ValueError("Incomplete varint")

def main():
    # Set up argument parsing
    parser = argparse.ArgumentParser(description='Deobfuscate data using common key and session ID.')
    parser.add_argument('--common-key', required=True, help='Common key string from setCommonKey.')
    parser.add_argument('--session-id', required=True, help='Session ID string from updateSessionId.')
    parser.add_argument('input_file', help='Path to the obfuscated data file.')
    
    args = parser.parse_args()
    
    # Build the XOR table
    xor_table = build_xor_table(args.common_key, args.session_id)
    
    if not xor_table:
        print("Error: XOR table is empty. Check the common key and session ID inputs.")
        sys.exit(1)
    
    # Read the obfuscated data
    try:
        with open(args.input_file, 'rb') as f:
            obfuscated_data = f.read()
    except Exception as e:
        print(f"Error reading input file: {e}")
        sys.exit(1)
    
    # Apply the conditional XOR operation
    data_after_xor = conditional_xor(obfuscated_data, xor_table)
    
    # Decode varint
    try:
        varint, varint_length = decode_varint(data_after_xor)
    except ValueError as e:
        print(f"Error decoding varint: {e}")
        sys.exit(1)
    
    print(f"=== Varint (Decompressed Size) ===")
    print(varint)
    print(f"Bytes consumed for varint: {varint_length}")
    
    # Extract compressed data
    compressed_data = data_after_xor[varint_length:]
    
    # Decompress using LZ4
    try:
        decompressed_data = lz4.block.decompress(compressed_data, uncompressed_size=varint)
    except lz4.block.LZ4BlockError as e:
        print(f"Error during LZ4 decompression: {e}")
        # Optionally, print the compressed data for inspection
        print("\nCompressed Data (Hex):")
        print(compressed_data.hex())
        sys.exit(1)
    
    # Print the decompressed data
    print("\n=== Decompressed Data ===")
    # Attempt to decode as UTF-8; if fails, print as hex
    try:
        print(decompressed_data.decode('utf-8'))
    except UnicodeDecodeError:
        print(decompressed_data.hex())

if __name__ == "__main__":
    main()