Affine Permutation For Obfuscation

aperm.py

#!/usr/bin/env python3
"""
Affine Permutation File Encoder/Decoder

This script allows you to encode and decode files using an Affine Permutation.
It uses command-line arguments to specify the operation mode (encode or decode), input file, and output file.
Additionally, it stores a SHA256 hash of the original data to verify successful decoding.

Usage:
    To encode a file:
        python perm.py -e input_file output_file

    To decode a file:
        python perm.py -d input_file output_file
        
    @modexpblog, October 2024
    
"""

import argparse
import struct
import random
import math
import sys
import hashlib
from pathlib import Path

# Updated Header Format: Magic (3s), a (I), c (I), SHA256 Hash (32s)
HEADER_FORMAT = '>3sII32s'  # Magic (3s), a (I), c (I), SHA256 (32s)
HEADER_SIZE = struct.calcsize(HEADER_FORMAT)
MAGIC = b'AP1'

def generate_parameters(m):
    """
    Generate suitable parameters 'a' and 'c' for Affine Permutation.
    'a' must be coprime with m, and 'c' must be non-zero.
    """
    if m == 0:
        raise ValueError("File is empty, cannot encode.")

    # Generate 'a' such that gcd(a, m) == 1
    while True:
        a = random.randint(1, m - 1)
        if math.gcd(a, m) == 1:
            break

    # Ensure 'c' is non-zero
    while True:
        c = random.randint(1, m-1)
        if c != 0:
            break

    return a, c

def affine_permutation(a, c, m):
    """
    Generate a permutation of indices [1, m-1] using Affine Permutation.
    """
    perm = []
    seen = set()
    for i in range(m):
        fi = (a * i + c) % m
        if fi in seen:
            raise ValueError("AP failed to generate a full permutation.")
        perm.append(fi)
        seen.add(fi)
    return perm

def shuffle_file(input_path, output_path):
    """
    Encode the input file and write the encoded data with header to the output file.
    """
    with open(input_path, 'rb') as f:
        data = f.read()

    m = len(data)
    if m == 0:
        print("Input file is empty. Nothing to encode.")
        sys.exit(1)

    a, c = generate_parameters(m)
    print(f"Encoding with parameters a={a}, c={c}, m={m}")

    # Compute SHA256 hash of original data
    sha256_hash = hashlib.sha256(data).digest()

    perm = affine_permutation(a, c, m)
    shuffled_data = bytes([data[p] for p in perm])
    
    # Pack header with SHA256 hash
    header = struct.pack(HEADER_FORMAT, MAGIC, a, c, sha256_hash)

    with open(output_path, 'wb') as f:
        f.write(header + shuffled_data)

    print(f"File encoded successfully. Output written to '{output_path}'.")

def unshuffle_file(input_path, output_path):
    """
    Decode the input file using the header parameters and write the original data to the output file.
    Verifies the SHA256 hash to ensure successful decoding.
    """
    with open(input_path, 'rb') as f:
        header = f.read(HEADER_SIZE)
        if len(header) < HEADER_SIZE:
            print("Input file is too small to contain a valid header.")
            sys.exit(1)

        magic, a, c, stored_hash = struct.unpack(HEADER_FORMAT, header)
        
        if magic != MAGIC:
            print("Invalid file format or missing header.")
            sys.exit(1)

        shuffled_data = f.read()

    m = len(shuffled_data)
    
    if m == 0:
        print("Encoded data is empty. Nothing to decode.")
        sys.exit(1)

    print(f"Decoding with parameters a={a}, c={c}, m={m}")

    perm = affine_permutation(a, c, m)

    # Create inverse permutation
    inverse_perm = [0] * m
    for i, p in enumerate(perm):
        inverse_perm[p] = i

    original_data = bytes([shuffled_data[inverse_perm[i]] for i in range(m)])

    # Compute SHA256 hash of reconstructed data
    reconstructed_hash = hashlib.sha256(original_data).digest()

    # Verify hash
    if reconstructed_hash != stored_hash:
        print("SHA256 hash mismatch! Decoding may have failed or data is corrupted.")
        sys.exit(1)
    else:
        print("SHA256 hash verified successfully. Decoding is successful.")

    with open(output_path, 'wb') as f:
        f.write(original_data)

    print(f"File decoded successfully. Output written to '{output_path}'.")

def parse_arguments():
    """
    Parse command-line arguments.
    """
    parser = argparse.ArgumentParser(description="Encode or decode a file using Affine Permutation with SHA256 verification.")
    
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-e', '--encode', action='store_true', help="Encode (shuffle) the input file.")
    group.add_argument('-d', '--decode', action='store_true', help="Decode (unshuffle) the input file.")

    parser.add_argument('input_file', help="Path to the input file.")
    parser.add_argument('output_file', help="Path to the output file.")

    return parser.parse_args()

def main():
    args = parse_arguments()

    input_path = Path(args.input_file)
    output_path = Path(args.output_file)

    if not input_path.is_file():
        print(f"Input file '{input_path}' does not exist or is not a file.")
        sys.exit(1)

    if output_path.exists():
        response = input(f"Output file '{output_path}' already exists. Overwrite? (y/n): ").lower()
        if response != 'y':
            print("Operation cancelled by user.")
            sys.exit(0)

    if args.encode:
        shuffle_file(input_path, output_path)
    elif args.decode:
        unshuffle_file(input_path, output_path)
    else:
        print("Invalid operation. Use '-e' to encode or '-d' to decode.")
        sys.exit(1)

if __name__ == '__main__':
    main()