#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import struct

def calculate_checksum(data):
    """计算小端序校验和：每4字节相加后取反（32位无符号）"""
    if len(data) % 4 != 0:
        raise ValueError("Data length must be multiple of 4")
    
    sum_val = 0
    # 遍历所有4字节块（最后4字节除外）
    for i in range(0, len(data)-4, 4):
        chunk = data[i:i+4]
        value = struct.unpack('<I', chunk)[0]  # 小端序解析
        sum_val = (sum_val + value) & 0xFFFFFFFF  # 32位溢出保护
    
    return (~sum_val) & 0xFFFFFFFF  # 取反并保持32位

def process_file(filename):
    """处理文件的主函数"""
    try:
        # 读取文件内容（二进制模式）
        with open(filename, 'rb+') as f:
            original_data = f.read()
            data_len = len(original_data)
            
            # ========== 验证阶段 ==========
            # 1. 检查文件长度是否为4的倍数
            if data_len % 4 != 0:
                raise ValueError(f"文件长度 {data_len} 不是4的倍数")
            
            # 2. 检查第二个4字节是否等于文件长度
            if data_len >= 8:
                declared_length = struct.unpack('<I', original_data[4:8])[0]
                if declared_length != data_len:
                    raise ValueError(
                        f"长度声明不匹配：文件头声明={declared_length}，实际={data_len}"
                    )
            
            # ========== 计算校验和 ==========
            checksum = calculate_checksum(original_data)
            
            # ========== 写入校验和 ==========
            # 定位到文件最后4字节
            f.seek(-4, os.SEEK_END)
            f.write(struct.pack('<I', checksum))  # 小端序写入
            
        print(f"[成功] 文件 {filename} 已更新校验码：0x{checksum:08X}")
        return True
        
    except Exception as e:
        print(f"[错误] 处理文件 {filename} 失败：{str(e)}")
        return False

if __name__ == "__main__":
    # 参数检查
    if len(sys.argv) != 2:
        print("用法: python checksum.py <文件名>")
        sys.exit(1)
    
    input_file = sys.argv[1]
    
    # 文件存在性检查
    if not os.path.isfile(input_file):
        print(f"[错误] 文件不存在: {input_file}")
        sys.exit(1)
    
    # 执行处理
    if not process_file(input_file):
        sys.exit(1)