Files
aegis-python/examples/benchmark.py
T

106 lines
3.0 KiB
Python

#!/usr/bin/env python3
"""
Python benchmark matching src/test/benchmark.zig for all supported Aegis algorithms.
It performs two benchmarks with the same parameters as the Zig version:
- AEGIS encrypt (attached tag, maclen = ABYTES_MIN)
- AEGIS MAC (clone state pattern)
Output format and throughput units mirror the Zig benchmark (Mb/s).
"""
import os
import time
from pyaegis import aegis128l, aegis128x2, aegis128x4, aegis256, aegis256x2, aegis256x4
MSG_LEN = 16384000 # 16 MiB
ITERATIONS = 100
ALGORITHMS = [
("AEGIS-128L", aegis128l),
("AEGIS-128X2", aegis128x2),
("AEGIS-128X4", aegis128x4),
("AEGIS-256", aegis256),
("AEGIS-256X2", aegis256x2),
("AEGIS-256X4", aegis256x4),
]
def _random_bytes(n: int) -> bytes:
return os.urandom(n)
def bench_encrypt(alg_name: str, a) -> None:
key = _random_bytes(a.KEYBYTES)
nonce = _random_bytes(a.NPUBBYTES)
# Single buffer, as in Zig: c_out == m buffer, with tag appended
maclen = a.ABYTES_MIN
buf = bytearray(MSG_LEN + maclen)
# Initialize buffer with random data
buf[:] = _random_bytes(len(buf))
mview = memoryview(buf)[:MSG_LEN]
t0 = time.perf_counter()
for _ in range(ITERATIONS):
a.encrypt(key, nonce, mview, None, maclen=maclen, into=buf)
t1 = time.perf_counter()
# Prevent any unrealistic optimization assumptions
_ = buf[0]
bits = MSG_LEN * ITERATIONS * 8
elapsed_s = t1 - t0
throughput_mbps = (
(bits / (elapsed_s * 1_000_000)) if elapsed_s > 0 else float("inf")
)
print(f"{alg_name}\t{throughput_mbps:10.2f} Mb/s")
def bench_mac(alg_name: str, a) -> None:
key = _random_bytes(a.KEYBYTES)
nonce = _random_bytes(a.NPUBBYTES)
buf = bytearray(MSG_LEN)
buf[:] = _random_bytes(len(buf))
mac0 = a.Mac(key, nonce)
mac_out = bytearray(a.ABYTES_MAX)
t0 = time.perf_counter()
for _ in range(ITERATIONS):
mac = mac0.clone()
mac.update(buf)
mac.final(maclen=a.ABYTES_MAX, into=mac_out)
t1 = time.perf_counter()
_ = mac_out[0]
bits = MSG_LEN * ITERATIONS * 8
elapsed_s = t1 - t0
throughput_mbps = (
(bits / (elapsed_s * 1_000_000)) if elapsed_s > 0 else float("inf")
)
print(f"{alg_name} MAC\t{throughput_mbps:10.2f} Mb/s")
if __name__ == "__main__":
# aegis_init() is called in the loader at import time already
# Run encrypt benchmarks in order: 256, 256x2, 256x4, 128l, 128x2, 128x4
bench_encrypt("AEGIS-256", aegis256)
bench_encrypt("AEGIS-256X2", aegis256x2)
bench_encrypt("AEGIS-256X4", aegis256x4)
bench_encrypt("AEGIS-128L", aegis128l)
bench_encrypt("AEGIS-128X2", aegis128x2)
bench_encrypt("AEGIS-128X4", aegis128x4)
# Run MAC benchmarks in order: 128l, 128x2, 128x4, 256, 256x2, 256x4
bench_mac("AEGIS-128L", aegis128l)
bench_mac("AEGIS-128X2", aegis128x2)
bench_mac("AEGIS-128X4", aegis128x4)
bench_mac("AEGIS-256", aegis256)
bench_mac("AEGIS-256X2", aegis256x2)
bench_mac("AEGIS-256X4", aegis256x4)