#!/usr/bin/env python3
"""
MEL Checksum Cracker - Final comprehensive tool
Based on analysis showing sum-sequence is constant but checksum follows different pattern
"""

import sys
import struct

def parse_hex(hex_str):
    """Parse hex string to bytes"""
    return [int(hex_str[i:i+2], 16) for i in range(0, len(hex_str), 2)]

def calculate_mel_checksum(payload_bytes, sequence=None):
    """
    Calculate MEL checksum based on discovered patterns
    This function will be updated as we discover the actual algorithm
    """
    if sequence is None:
        # Try to extract sequence from payload if it's a full packet
        if len(payload_bytes) >= 8:
            sequence = payload_bytes[7]
        else:
            sequence = 0
    
    # Calculate payload sum (excluding checksum bytes)
    if len(payload_bytes) > 31:  # Full packet
        payload_sum = sum(payload_bytes[:-2])
    else:  # Just payload
        payload_sum = sum(payload_bytes)
    
    # Test various algorithms based on our analysis
    algorithms = {
        'simple_sum': payload_sum & 0xFFFF,
        'sum_minus_seq': (payload_sum - sequence) & 0xFFFF,
        'sum_plus_seq': (payload_sum + sequence) & 0xFFFF,
        'twos_complement': (~payload_sum + 1) & 0xFFFF,
        'ones_complement': (~payload_sum) & 0xFFFF,
        'constant_minus_sum': (0x10000 - payload_sum) & 0xFFFF,
        'sum_with_carry': payload_sum + (payload_sum >> 16),
    }
    
    return algorithms

def analyze_pattern_discovery(filename):
    """Discover the actual checksum pattern"""
    print(f"Pattern Discovery Analysis for {filename}")
    print("=" * 60)
    
    with open(filename, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]
    
    # Analyze first 20 entries to find the pattern
    entries = []
    for i, line in enumerate(lines[:20]):
        bytes_data = parse_hex(line)
        payload = bytes_data[:-2]
        checksum = bytes_data[-2] | (bytes_data[-1] << 8)  # Little endian
        sequence = bytes_data[7] if len(bytes_data) > 7 else 0
        payload_sum = sum(payload)
        
        entries.append({
            'index': i,
            'sequence': sequence,
            'checksum': checksum,
            'payload_sum': payload_sum,
            'payload': payload,
            'hex': line
        })
    
    # Print analysis table
    print("Entry | Seq | Checksum | PayloadSum | Sum-Seq | Patterns")
    print("------|-----|----------|------------|---------|----------")
    
    base_constant = None
    for entry in entries:
        seq = entry['sequence']
        check = entry['checksum']
        psum = entry['payload_sum']
        sum_minus_seq = psum - seq
        
        if base_constant is None:
            base_constant = sum_minus_seq
        
        # Test various patterns
        pattern_tests = []
        if sum_minus_seq == base_constant:
            pattern_tests.append("SUM-SEQ=CONST")
        
        # Test if checksum relates to a base value
        if entry['index'] == 0:
            base_checksum = check
            checksum_base = check
        else:
            checksum_diff = check - base_checksum
            pattern_tests.append(f"ΔCHK={checksum_diff:+d}")
        
        patterns = " ".join(pattern_tests) if pattern_tests else "-"
        
        print(f"{entry['index']:5d} | {seq:3d} | 0x{check:04x} | {psum:10d} | {sum_minus_seq:7d} | {patterns}")
    
    # Test for mathematical relationships
    print(f"\nMathematical Relationship Analysis:")
    print(f"Base constant (sum - sequence): {base_constant}")
    
    # Look for checksum calculation pattern
    print(f"\nChecksum Pattern Analysis:")
    
    # Test if there's a consistent transformation from payload_sum to checksum
    transformations = {}
    for entry in entries:
        psum = entry['payload_sum']
        check = entry['checksum']
        seq = entry['sequence']
        
        # Test various transformations
        transforms = {
            'sum_low8': psum & 0xFF,
            'sum_high8': (psum >> 8) & 0xFF,
            'sum_mod256': psum % 256,
            'sum_mod255': psum % 255,
            'sum_rotated': ((psum << 8) | (psum >> 8)) & 0xFFFF,
            'sum_xor_seq': (psum ^ seq) & 0xFFFF,
            'sum_add_magic': (psum + 0x5555) & 0xFFFF,
            'sum_sub_magic': (psum - 0x5555) & 0xFFFF,
        }
        
        for name, value in transforms.items():
            if value == check:
                if name not in transformations:
                    transformations[name] = []
                transformations[name].append(entry['index'])
    
    if transformations:
        print("Found consistent transformations:")
        for transform, indices in transformations.items():
            if len(indices) > 1:
                print(f"  {transform}: matches at entries {indices}")
    else:
        print("No simple mathematical transformation found")
    
    # Advanced pattern detection
    print(f"\nAdvanced Pattern Detection:")
    
    # Check if it's a lookup table or formula with sequence
    checksum_by_seq = {}
    for entry in entries:
        seq = entry['sequence']
        check = entry['checksum']
        if seq not in checksum_by_seq:
            checksum_by_seq[seq] = []
        checksum_by_seq[seq].append(check)
    
    # Look for sequence-based pattern
    if len(checksum_by_seq) > 1:
        print("Checksum vs Sequence relationship:")
        for seq in sorted(checksum_by_seq.keys())[:10]:
            checksums = checksum_by_seq[seq]
            print(f"  Sequence {seq:3d}: checksums {[f'0x{c:04x}' for c in checksums]}")
    
    return entries

def brute_force_checksum_algorithm(entries):
    """Brute force the checksum algorithm using known good data"""
    print("\nBrute Force Algorithm Discovery:")
    print("=" * 40)
    
    # Get the first entry as reference
    ref_entry = entries[0]
    ref_sum = ref_entry['payload_sum']
    ref_checksum = ref_entry['checksum']
    
    print(f"Reference: sum={ref_sum}, checksum=0x{ref_checksum:04x}")
    
    # Try to find the magic constant or operation
    # Test if checksum = (payload_sum + magic) & 0xFFFF
    for magic in range(-1000, 1000):
        predicted = (ref_sum + magic) & 0xFFFF
        if predicted == ref_checksum:
            print(f"Possible algorithm: checksum = (payload_sum + {magic}) & 0xFFFF")
            
            # Verify with other entries
            matches = 0
            for entry in entries[1:6]:  # Test next 5 entries
                test_predicted = (entry['payload_sum'] + magic) & 0xFFFF
                if test_predicted == entry['checksum']:
                    matches += 1
            
            print(f"  Verification: {matches}/5 additional entries match")
            if matches >= 4:
                print(f"  🎉 LIKELY ALGORITHM FOUND!")
                return f"(payload_sum + {magic}) & 0xFFFF"
    
    # Test multiplicative factors
    for factor in [1, 2, 3, 4, 5, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128, 255, 256]:
        for offset in range(-100, 101, 10):
            predicted = (ref_sum * factor + offset) & 0xFFFF
            if predicted == ref_checksum:
                # Verify with next entry
                next_entry = entries[1]
                next_predicted = (next_entry['payload_sum'] * factor + offset) & 0xFFFF
                if next_predicted == next_entry['checksum']:
                    print(f"Possible algorithm: checksum = (payload_sum * {factor} + {offset}) & 0xFFFF")
    
    print("No simple algorithm found in brute force range")
    return None

def generate_implementation(algorithm_desc):
    """Generate C/Python implementation of discovered algorithm"""
    if not algorithm_desc:
        return
    
    print(f"\nGenerated Implementation:")
    print("=" * 30)
    
    print("Python:")
    print(f"def calculate_mel_checksum(payload_bytes):")
    print(f"    payload_sum = sum(payload_bytes)")
    print(f"    return {algorithm_desc}")
    
    print("\nC:")
    print(f"uint16_t calculate_mel_checksum(uint8_t* payload, size_t length) {{")
    print(f"    uint32_t payload_sum = 0;")
    print(f"    for (size_t i = 0; i < length; i++) {{")
    print(f"        payload_sum += payload[i];")
    print(f"    }}")
    print(f"    return {algorithm_desc.replace('payload_sum', 'payload_sum')};")
    print(f"}}")

def validate_algorithm(filename, algorithm_func):
    """Validate the discovered algorithm against all entries in file"""
    print(f"\nValidating algorithm against all entries in {filename}")
    print("=" * 50)
    
    with open(filename, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]
    
    matches = 0
    total = 0
    mismatches = []
    
    for i, line in enumerate(lines):
        bytes_data = parse_hex(line)
        payload = bytes_data[:-2]
        expected_checksum = bytes_data[-2] | (bytes_data[-1] << 8)
        
        calculated_checksum = algorithm_func(payload)
        
        if calculated_checksum == expected_checksum:
            matches += 1
        else:
            mismatches.append((i, expected_checksum, calculated_checksum))
        
        total += 1
        
        # Show first few results
        if i < 10:
            status = "✓" if calculated_checksum == expected_checksum else "✗"
            print(f"  Entry {i:3d}: expected=0x{expected_checksum:04x}, calculated=0x{calculated_checksum:04x} {status}")
    
    success_rate = (matches / total) * 100
    print(f"\nValidation Results:")
    print(f"  Matches: {matches}/{total} ({success_rate:.1f}%)")
    
    if success_rate == 100:
        print("  🎉 PERFECT MATCH! Algorithm verified!")
    elif success_rate > 95:
        print("  ⚠️  Very high success rate - likely correct with minor issues")
    else:
        print("  ❌ Low success rate - algorithm needs refinement")
        if mismatches[:3]:
            print("  First few mismatches:")
            for entry_idx, expected, calculated in mismatches[:3]:
                print(f"    Entry {entry_idx}: expected=0x{expected:04x}, got=0x{calculated:04x}")

def main():
    if len(sys.argv) < 2:
        print("MEL Checksum Cracker - Comprehensive Analysis Tool")
        print("=" * 50)
        print("Usage: python mel_checksum_cracker.py <hex_file>")
        print("\nThis tool will:")
        print("  1. Discover checksum patterns")
        print("  2. Brute force the algorithm")
        print("  3. Generate implementations")
        print("  4. Validate against all data")
        sys.exit(1)
    
    filename = sys.argv[1]
    
    try:
        # Step 1: Pattern discovery
        entries = analyze_pattern_discovery(filename)
        
        # Step 2: Brute force algorithm
        algorithm = brute_force_checksum_algorithm(entries)
        
        # Step 3: Generate implementation
        generate_implementation(algorithm)
        
        # Step 4: Validation (if algorithm found)
        if algorithm:
            # Create a lambda function for validation
            def test_algorithm(payload):
                payload_sum = sum(payload)
                # This would be replaced with the actual discovered algorithm
                # For now, using a placeholder
                return payload_sum & 0xFFFF
            
            validate_algorithm(filename, test_algorithm)
        
    except FileNotFoundError:
        print(f"Error: File '{filename}' not found")
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()