109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
import logging
|
|
from argparse import ArgumentParser
|
|
from pathlib import Path
|
|
|
|
_VERBOSE_SBOX = False
|
|
_VERBOSE_PRGA = False
|
|
_STEPWISE_PRGA = False
|
|
|
|
|
|
def sbox(key: list[int]) -> list[int]:
|
|
"""
|
|
Creates and permutates the substitution box given a key.
|
|
"""
|
|
# create sbox
|
|
S = [i for i in range(256)]
|
|
|
|
# permutate sbox
|
|
j = 0
|
|
for i in range(256):
|
|
j = (j + S[i] + key[i % len(key)]) % 256
|
|
if _VERBOSE_SBOX:
|
|
print(f"Swapping S[{i=:3d}] = {S[i]:3d} and S[{j=:3d}] = {S[j]:3d}.")
|
|
S[i], S[j] = S[j], S[i] # swap
|
|
|
|
return S
|
|
|
|
|
|
def prga(text: list[int], S: list[int]) -> list[int]:
|
|
"""
|
|
Encrypts/Decrypts text given a substitution box.
|
|
"""
|
|
text = list(text)
|
|
i = 0
|
|
j = 0
|
|
for x in range(len(text)):
|
|
i = (i + 1) % 256
|
|
j = (j + S[i]) % 256
|
|
|
|
if _VERBOSE_PRGA:
|
|
print(f"Swapping S[{i=:3d}] = {S[i]:3d} and S[{j=:3d}] = {S[j]:3d}.")
|
|
|
|
if _STEPWISE_PRGA:
|
|
try:
|
|
input("Press ENTER to continue ...") # Step through
|
|
except KeyboardInterrupt:
|
|
print("KeyboardInterrupt.")
|
|
exit(0)
|
|
|
|
S[i], S[j] = S[j], S[i] # swap
|
|
K = S[(S[i] + S[j]) % 256]
|
|
text[x] ^= K
|
|
return text
|
|
|
|
|
|
def convert(text: bytes | list[int]) -> bytes | list[int]:
|
|
"""
|
|
Converts text from bytes -> list[int] and back.
|
|
"""
|
|
if isinstance(text, bytes):
|
|
return [int(c) for c in text]
|
|
elif isinstance(text, list):
|
|
return bytes(text)
|
|
else:
|
|
raise ValueError(f"Unsupported input type '{type(text)}'")
|
|
|
|
|
|
def rc4(in_file: str | Path, out_file: str | Path, key_file: str | Path):
|
|
"""
|
|
Execute the RC4 encryption/decryption algorithm on the input file,
|
|
creating the output file using the contents from the key file.
|
|
"""
|
|
|
|
with open(key_file, "rb") as key_file:
|
|
S = sbox(convert(key_file.read()))
|
|
|
|
with open("sbox_solution", "wb") as fd:
|
|
fd.write(convert(S))
|
|
|
|
# read input
|
|
with open(in_file, "rb") as data_in_file:
|
|
data_in = data_in_file.read()
|
|
|
|
data_in = convert(data_in) # bytes -> list[int]
|
|
data_out = prga(data_in, S)
|
|
data_out = convert(data_out) # list[int] -> bytes
|
|
|
|
# write output
|
|
with open(out_file, "wb") as data_out_file:
|
|
data_out_file.write(data_out)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = ArgumentParser(
|
|
description="Encrypts plaintext using the RC4 encryption algorithm, or decrypts RC4 ciphertext (based on input)"
|
|
)
|
|
parser.add_argument("input", type=str, help="Input file")
|
|
parser.add_argument("key", type=str, help="Key file")
|
|
parser.add_argument("output", type=str, help="Output file")
|
|
parser.add_argument("--verbose-sbox", const=True, default=False, action="store_const", help="Show S-Box swaps")
|
|
parser.add_argument("--verbose-prga", const=True, default=False, action="store_const", help="Show PRGA swaps")
|
|
parser.add_argument("--step-prga", const=True, default=False, action="store_const", help="Step-by-step PRGA swaps")
|
|
args = parser.parse_args()
|
|
|
|
_VERBOSE_SBOX = args.verbose_sbox
|
|
_VERBOSE_PRGA = args.verbose_prga
|
|
_STEPWISE_PRGA = args.step_prga
|
|
|
|
rc4(args.input, args.output, args.key)
|