benchmark.py - master - wspppe/AES

benchmark.py
138 строк5 KB
wspppe
create: benchmark.py
17 май 2026, 14:16
17 май 2026, 14:16a366189
import os
import time
import subprocess
from Crypto.Cipher import AES

FILE_SIZES_MB = [1, 10, 100, 500, 1000]  
KEY = b"0011223344556677"        
IV = b"0001020304050607"           
NONCE = b"000102030405"           
CHUNK_SIZE = 64 * 1024             

def generate_file(name, size_mb):
    if os.path.exists(name):
        print(f"Файл {name} уже существует, пропускаем генерацию.")
        return
    print(f"Генерация файла {name} ({size_mb} MB)")
    with open(name, "wb") as f:
        for _ in range((size_mb * 1024 * 1024) // CHUNK_SIZE):
            f.write(os.urandom(CHUNK_SIZE))

def benchmark_openssl(mode, infile, outfile, decryptfile, is_decrypt=False):
    key_hex = KEY.hex()
    iv_hex = IV.hex()
    
    cmd = ["openssl", "enc", f"-aes-128-{mode}", "-K", key_hex, "-iv", iv_hex]
    if is_decrypt:
        cmd += ["-d", "-in", outfile, "-out", decryptfile]
    else:
        cmd += ["-e", "-in", infile, "-out", outfile]
        
    start = time.perf_counter()
    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
    end = time.perf_counter()
    return end - start

def benchmark_gcm_large(infile, outfile, decryptfile, is_decrypt=False):
    """Потоковое шифрование/расшифрование GCM для больших файлов"""
    start = time.perf_counter()
    cipher = AES.new(KEY, AES.MODE_GCM, nonce=NONCE)
    
    if not is_decrypt:
        with open(infile, "rb") as f_in, open(outfile, "wb") as f_out:
            while True:
                chunk = f_in.read(CHUNK_SIZE)
                if not chunk:
                    break
                f_out.write(cipher.encrypt(chunk))
            f_out.write(cipher.digest())
    else:
        with open(outfile, "rb") as f_in, open(decryptfile, "wb") as f_out:
            file_size = os.path.getsize(outfile)
            if file_size < 16:
                f_out.write(b"ERROR: File too small")
                return time.perf_counter() - start
                
            ciphertext_len = file_size - 16
            
            bytes_read = 0
            while bytes_read < ciphertext_len:
                to_read = min(CHUNK_SIZE, ciphertext_len - bytes_read)
                chunk = f_in.read(to_read)
                if not chunk:
                    break
                f_out.write(cipher.decrypt(chunk))
                bytes_read += len(chunk)

            tag = f_in.read(16)
            try:
                cipher.verify(tag)
            except ValueError:
                f_out.seek(0)
                f_out.truncate()
                f_out.write(b"ERROR: Integrity check failed!")
                
    end = time.perf_counter()
    return end - start

def inject_bit_error(filename):
    with open(filename, "r+b") as f:
        f.seek(os.path.getsize(filename) // 2)
        byte = f.read(1)
        if byte:
            f.seek(-1, 1)
            f.write(bytes([byte[0] ^ 1]))

print("Старт исследования")
results = []

for size in FILE_SIZES_MB:
    infile = f"test_{size}M.bin"
    generate_file(infile, size)
    
    for mode in ["cbc", "ctr", "gcm"]:
        outfile = f"enc_{mode}_{size}M.bin"
        decfile = f"dec_{mode}_{size}M.bin"
        
        if mode in ["cbc", "ctr"]:
            t_enc = benchmark_openssl(mode, infile, outfile, decfile, is_decrypt=False)
            t_dec = benchmark_openssl(mode, infile, outfile, decfile, is_decrypt=True)
        else:
            t_enc = benchmark_gcm_large(infile, outfile, decfile, is_decrypt=False)
            t_dec = benchmark_gcm_large(infile, outfile, decfile, is_decrypt=True)
            
        print(f"Режим {mode.upper()} [{size}MB] -> Шифрование: {t_enc:.2f}с | Дешифрование: {t_dec:.2f}с")
        results.append(f"{size}MB;{mode};{t_enc:.4f};{t_dec:.4f}\n")

with open("benchmark_results.csv", "w") as f:
    f.write("Size;Mode;EncTime;DecTime\n")
    f.writelines(results)

print("\n Тест на ошибку (На базе файла 1 МБ)")
for mode in ["cbc", "ctr", "gcm"]:
    target_file = f"enc_{mode}_1M.bin"
    corrupted_dec_file = f"dec_{mode}_1M_corrupted.bin"
    
    inject_bit_error(target_file)
    
    if mode in ["cbc", "ctr"]:
        try:
            benchmark_openssl(mode, "test_1M.bin", target_file, corrupted_dec_file, is_decrypt=True)
            print(f"Режим {mode.upper()}: Расшифрован с искажением.")
        except Exception:
            print(f"Режим {mode.upper()}: Ошибка OpenSSL.")
    else:
        benchmark_gcm_large("test_1M.bin", target_file, corrupted_dec_file, is_decrypt=True)
        with open(corrupted_dec_file, "rb") as f:
            print(f"Режим GCM: {f.read(50).decode('utf-8', errors='ignore')}")

print("\nГотово.")
for size in FILE_SIZES_MB:
    if size > 1:
        for mode in ["cbc", "ctr", "gcm"]:
            try:
                os.remove(f"test_{size}M.bin")
                os.remove(f"enc_{mode}_{size}M.bin")
                os.remove(f"dec_{mode}_{size}M.bin")
            except FileNotFoundError:
                pass