#!/usr/bin/env python3
"""
Comprehensive checksum bruteforce tool for the MEL protocol
Usage: python checksum_bruteforce.py <hex_file>
"""

import sys
import struct
from typing import List, Tuple, Dict

def parse_hex_line(hex_string: str) -> List[int]:
    """Parse a hex string into list of bytes"""
    hex_clean = hex_string.strip().replace(' ', '')
    return [int(hex_clean[i:i+2], 16) for i in range(0, len(hex_clean), 2)]

def bytes_to_hex(bytes_list: List[int]) -> str:
    """Convert bytes to hex string"""
    return ''.join(f'{b:02x}' for b in bytes_list)

class ChecksumTester:
    def __init__(self):
        self.algorithms = [
            self.simple_sum,
            self.sum_with_carry,
            self.twos_complement,
            self.ones_complement,
            self.xor_checksum,
            self.crc16_ccitt,
            self.crc16_ibm,
            self.fletcher16,
            self.modsum_256,
            self.internet_checksum,
        ]
        
    def simple_sum(self, data: List[int]) -> int:
        """Simple sum of all bytes"""
        return sum(data) & 0xFFFF
    
    def sum_with_carry(self, data: List[int]) -> int:
        """Sum with end-around carry"""
        s = sum(data)
        while s > 0xFFFF:
            s = (s & 0xFFFF) + (s >> 16)
        return s
    
    def twos_complement(self, data: List[int]) -> int:
        """Two's complement of sum"""
        s = sum(data)
        return (~s + 1) & 0xFFFF
    
    def ones_complement(self, data: List[int]) -> int:
        """One's complement of sum"""
        s = sum(data)
        return (~s) & 0xFFFF
    
    def xor_checksum(self, data: List[int]) -> int:
        """XOR of all bytes, extended to 16-bit"""
        result = 0
        for b in data:
            result ^= b
        return result
    
    def crc16_ccitt(self, data: List[int], poly: int = 0x1021) -> int:
        """CRC-16 CCITT"""
        crc = 0xFFFF
        for byte in data:
            crc ^= (byte << 8)
            for _ in range(8):
                if crc & 0x8000:
                    crc = (crc << 1) ^ poly
                else:
                    crc <<= 1
                crc &= 0xFFFF
        return crc
    
    def crc16_ibm(self, data: List[int]) -> int:
        """CRC-16 IBM/ANSI"""
        return self.crc16_ccitt(data, 0x8005)
    
    def fletcher16(self, data: List[int]) -> int:
        """Fletcher-16 checksum"""
        sum1 = sum2 = 0
        for byte in data:
            sum1 = (sum1 + byte) % 255
            sum2 = (sum2 + sum1) % 255
        return (sum2 << 8) | sum1
    
    def modsum_256(self, data: List[int]) -> int:
        """Sum modulo 256, extended to 16-bit"""
        return sum(data) % 256
    
    def internet_checksum(self, data: List[int]) -> int:
        """Internet/TCP checksum"""
        # Pad to even length
        if len(data) % 2:
            data = data + [0]
        
        s = 0
        for i in range(0, len(data), 2):
            s += (data[i] << 8) + data[i+1]
        
        while s >> 16:
            s = (s & 0xFFFF) + (s >> 16)
        
        return (~s) & 0xFFFF

def test_parametric_algorithms(data: List[int], expected: int) -> List[str]:
    """Test algorithms with various parameters"""
    matches = []
    
    # Test sum with different initial values
    for init_val in range(0, 0x10000, 0x1000):
        result = (init_val + sum(data)) & 0xFFFF
        if result == expected:
            matches.append(f"Sum + 0x{init_val:04x}")
    
    # Test sum with different modulo values
    for mod_val in [0xFF, 0x100, 0x101, 0x1FF, 0x200, 0xFFFF, 0x10000]:
        if mod_val > 0:
            result = sum(data) % mod_val
            if result == expected:
                matches.append(f"Sum mod 0x{mod_val:x}")
    
    # Test XOR with different patterns
    for pattern in [0x00, 0xFF, 0xAA, 0x55, 0x5A, 0xA5]:
        result = 0
        for byte in data:
            result ^= (byte ^ pattern)
        if (result & 0xFFFF) == expected:
            matches.append(f"XOR with pattern 0x{pattern:02x}")
    
    # Test rotation-based checksums
    for shift in range(1, 16):
        result = 0
        for byte in data:
            result = ((result << shift) | (result >> (16 - shift))) & 0xFFFF
            result ^= byte
        if result == expected:
            matches.append(f"Rotate-XOR shift {shift}")
    
    return matches

def analyze_file(filename: str):
    """Analyze a file of hex data to find checksum patterns"""
    with open(filename, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]
    
    print(f"Analyzing {filename} with {len(lines)} entries")
    
    tester = ChecksumTester()
    algorithm_matches = {}
    
    # Test first 10 entries to find patterns
    for i, line in enumerate(lines[:10]):
        bytes_data = parse_hex_line(line)
        if len(bytes_data) < 3:
            continue
            
        # Assume last 2 bytes are checksum
        payload = bytes_data[:-2]
        checksum_be = (bytes_data[-2] << 8) | bytes_data[-1]
        checksum_le = bytes_data[-2] | (bytes_data[-1] << 8)
        
        print(f"\nEntry {i}:")
        print(f"  Payload: {bytes_to_hex(payload)}")
        print(f"  Checksum BE: 0x{checksum_be:04x}")
        print(f"  Checksum LE: 0x{checksum_le:04x}")
        
        # Test standard algorithms for both byte orders
        for checksum, order in [(checksum_be, "BE"), (checksum_le, "LE")]:
            print(f"  Testing {order} interpretation:")
            for algo in tester.algorithms:
                result = algo(payload)
                if result == checksum:
                    algo_name = f"{algo.__name__}_{order}"
                    algorithm_matches[algo_name] = algorithm_matches.get(algo_name, 0) + 1
                    print(f"    ✓ {algo.__name__}: 0x{result:04x}")
                else:
                    print(f"    ✗ {algo.__name__}: 0x{result:04x}")
            
            # Test parametric algorithms
            param_matches = test_parametric_algorithms(payload, checksum)
            for match in param_matches:
                print(f"    ✓ {match}")
                key = f"{match}_{order}"
                algorithm_matches[key] = algorithm_matches.get(key, 0) + 1
    
    # Summary of algorithms that work consistently
    print(f"\n{'='*50}")
    print("SUMMARY - Algorithms with multiple matches:")
    for algo, count in sorted(algorithm_matches.items(), key=lambda x: x[1], reverse=True):
        if count > 1:
            print(f"  {algo}: {count} matches")
    
    return algorithm_matches

def brute_force_unknown_algorithm(filename: str, max_entries: int = 5):
    """Brute force approach for completely unknown algorithms"""
    with open(filename, 'r') as f:
        lines = [line.strip() for line in f if line.strip()][:max_entries]
    
    print(f"Brute forcing {len(lines)} entries...")
    
    # Collect all data points
    data_points = []
    for line in lines:
        bytes_data = parse_hex_line(line)
        payload = bytes_data[:-2]
        checksum_le = bytes_data[-2] | (bytes_data[-1] << 8)
        data_points.append((payload, checksum_le))
    
    # Try to find a mathematical relationship
    # This is a simplified approach - in practice you'd want more sophisticated analysis
    
    # Check if it's a linear relationship: checksum = a * sum(payload) + b
    if len(data_points) >= 2:
        payload_sums = [sum(payload) for payload, _ in data_points]
        checksums = [checksum for _, checksum in data_points]
        
        print(f"Payload sums: {[f'0x{s:x}' for s in payload_sums[:5]]}")
        print(f"Checksums:    {[f'0x{c:04x}' for c in checksums[:5]]}")
        
        # Try to solve for linear relationship
        if len(set(payload_sums)) > 1:  # Need different sums to solve
            sum1, sum2 = payload_sums[0], payload_sums[1]
            check1, check2 = checksums[0], checksums[1]
            
            if sum1 != sum2:
                # Solve: check1 = a * sum1 + b, check2 = a * sum2 + b
                a = (check2 - check1) / (sum2 - sum1)
                b = check1 - a * sum1
                
                print(f"Testing linear relationship: checksum = {a:.3f} * sum + {b:.3f}")
                
                # Verify with all data points
                matches = 0
                for payload, expected in data_points:
                    predicted = int(a * sum(payload) + b) & 0xFFFF
                    if predicted == expected:
                        matches += 1
                        print(f"  ✓ Linear match: sum={sum(payload)}, expected=0x{expected:04x}, predicted=0x{predicted:04x}")
                    else:
                        print(f"  ✗ Linear miss: sum={sum(payload)}, expected=0x{expected:04x}, predicted=0x{predicted:04x}")
                
                if matches == len(data_points):
                    print(f"🎉 FOUND LINEAR RELATIONSHIP! checksum = {a:.3f} * sum(payload) + {b:.3f}")

def generate_test_vectors(base_hex: str, variations: int = 10):
    """Generate test vectors by modifying a base message"""
    base_bytes = parse_hex_line(base_hex)
    payload = base_bytes[:-2]
    
    print(f"Generating {variations} test vectors from base:")
    print(f"Base: {base_hex}")
    
    # Generate variations by changing single bytes
    for i in range(min(variations, len(payload))):
        modified = payload.copy()
        modified[i] = (modified[i] + 1) % 256  # Increment one byte
        
        # You would calculate the correct checksum here and append it
        # For now, we'll just show the modified payload
        print(f"Variation {i}: {bytes_to_hex(modified)} + [CHECKSUM_TO_CALCULATE]")

def test_specific_algorithms(data: List[int], expected: int) -> Dict[str, int]:
    """Test specific algorithms that might be used in embedded systems"""
    results = {}
    
    # Test sum with various bit operations
    s = sum(data)
    
    # Basic variations
    results["sum_low16"] = s & 0xFFFF
    results["sum_high16"] = (s >> 16) & 0xFFFF
    results["sum_rotated"] = ((s << 8) | (s >> 8)) & 0xFFFF
    results["sum_inverted"] = (~s) & 0xFFFF
    results["sum_twos_comp"] = (-s) & 0xFFFF
    
    # With different initial values (common in embedded systems)
    for init in [0x0000, 0x5555, 0xAAAA, 0xFFFF, 0x1234, 0x4321]:
        results[f"sum_init_{init:04x}"] = (s + init) & 0xFFFF
        results[f"xor_init_{init:04x}"] = (s ^ init) & 0xFFFF
    
    # Byte-wise operations
    byte_xor = 0
    byte_sum = 0
    for b in data:
        byte_xor ^= b
        byte_sum += b
    
    results["byte_xor"] = byte_xor
    results["byte_xor_16bit"] = (byte_xor << 8) | byte_xor
    results["byte_sum_mod256"] = byte_sum % 256
    results["byte_sum_mod255"] = byte_sum % 255
    
    # Position-weighted sums
    pos_sum = sum(i * b for i, b in enumerate(data))
    results["position_weighted"] = pos_sum & 0xFFFF
    
    # Polynomial checksums (simplified)
    poly_result = 0
    for b in data:
        poly_result = ((poly_result << 1) ^ b) & 0xFFFF
    results["poly_shift_xor"] = poly_result
    
    # Return only matches
    return {name: value for name, value in results.items() if value == expected}

def main():
    if len(sys.argv) != 2:
        print("Usage: python checksum_bruteforce.py <hex_file>")
        print("\nThis tool will:")
        print("1. Test standard checksum algorithms")
        print("2. Try parametric variations")
        print("3. Attempt to find mathematical relationships")
        print("4. Generate test vectors for validation")
        sys.exit(1)
    
    filename = sys.argv[1]
    
    try:
        # Main analysis
        print("=" * 60)
        print("COMPREHENSIVE CHECKSUM ANALYSIS")
        print("=" * 60)
        
        matches = analyze_file(filename)
        
        print("\n" + "=" * 60)
        print("BRUTE FORCE UNKNOWN ALGORITHMS")
        print("=" * 60)
        
        brute_force_unknown_algorithm(filename, max_entries=10)
        
        print("\n" + "=" * 60)
        print("SPECIFIC EMBEDDED ALGORITHMS")
        print("=" * 60)
        
        # Test first entry with specific algorithms
        with open(filename, 'r') as f:
            first_line = f.readline().strip()
            
        bytes_data = parse_hex_line(first_line)
        payload = bytes_data[:-2]
        checksum_le = bytes_data[-2] | (bytes_data[-1] << 8)
        
        specific_matches = test_specific_algorithms(payload, checksum_le)
        if specific_matches:
            print("Specific algorithm matches found:")
            for name, value in specific_matches.items():
                print(f"  ✓ {name}: 0x{value:04x}")
        else:
            print("No specific algorithm matches found")
        
        print("\n" + "=" * 60)
        print("TEST VECTOR GENERATION")
        print("=" * 60)
        
        generate_test_vectors(first_line, 5)
        
    except FileNotFoundError:
        print(f"Error: File '{filename}' not found")
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    main()