Checksum.py

import hashlib #MD-5 checksum
import os
import zlib # CRC-32 checksum (faster, less secure)
import yaml
import shutil

# Globals
# Define the file to store checksums and backup
CHECKSUMS_FILE = "CHECKSUMS.yaml"
BACKUP_FILE = CHECKSUMS_FILE + ".old_backup"

def compute_checksum_MD5(file_path):
    # Compute checksum for a file
    hasher = hashlib.md5()
    with open(file_path, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            hasher.update(chunk)
    return hasher.hexdigest()

def compute_checksum_CRC32(file_path, last_bytes_to_include = 21734636):
    # Compute CRC-32 checksum for the last specified bytes in a file
    checksum = 0
    with open(file_path, 'rb') as f:
        file_size = os.path.getsize(file_path)
        # Move the file pointer to the beginning of the last specified bytes or beginning if file size is smaller than specified bytes
        f.seek(max(-last_bytes_to_include, -file_size), os.SEEK_END)
        # Read and update checksum for the specified range
        while True:
            chunk = f.read(min(4096, last_bytes_to_include))
            if not chunk:
                break
            checksum = zlib.crc32(chunk, checksum) & 0xFFFFFFFF
    return format(checksum, '08x')  # Convert the checksum to a hexadecimal string

def store_checksum(file_path, checksum, key):
    # Create a backup before making changes if the file exists
    if os.path.exists(CHECKSUMS_FILE):
        shutil.copyfile(CHECKSUMS_FILE, BACKUP_FILE)
    else:
        None
    # Load existing checksums from file
    checksums = load_yaml(CHECKSUMS_FILE)
    # Add a new value to the specified key
    if key not in checksums:
        checksums[key] = {}

    checksums[key][file_path] = checksum
    # Store checksums back to the file
    save_yaml(checksums, CHECKSUMS_FILE)

def load_checksums():
    # Load existing checksums from file
    return load_yaml(CHECKSUMS_FILE)

def load_yaml(file_path):
    if not os.path.exists(file_path):
        return {}
    with open(file_path, 'r') as file:
        data = yaml.safe_load(file)
    return data or {}

def save_yaml(data, file_path):
    with open(file_path, 'w') as file:
        yaml.dump(data, file, default_flow_style=False)

def checksum_exists(key, checksum):
    checksums = load_checksums()
    return key in checksums and any(checksum == stored_checksum for stored_checksum in checksums.get(key, {}).values())


# Example usage:
if __name__ == "__main__":
    checksums = load_checksums()
    path = "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/10_Martini2_nonPW/2_identical/1nnx/Take_5/md_0_1Mu.xtc"
    checksum = compute_checksum_CRC32(path)
    store_checksum(path, checksum, "key_1")

    if checksum_exists("key_1", checksum):
        print(f"Checksum {checksum} exists under key_1")
    else:
        print(f"Checksum {checksum} does not exist under key_1")