rsa_cipher.py

importos
importsys

from . importrsa_key_generatorasrkg

DEFAULT_BLOCK_SIZE=128
BYTE_SIZE=256


defget_blocks_from_text(
message: str, block_size: int=DEFAULT_BLOCK_SIZE
) ->list[int]:
message_bytes=message.encode("ascii")

block_ints= []
forblock_startinrange(0, len(message_bytes), block_size):
block_int=0
foriinrange(block_start, min(block_start+block_size, len(message_bytes))):
block_int+=message_bytes[i] * (BYTE_SIZE** (i%block_size))
block_ints.append(block_int)
returnblock_ints


defget_text_from_blocks(
block_ints: list[int], message_length: int, block_size: int=DEFAULT_BLOCK_SIZE
) ->str:
message: list[str] = []
forblock_intinblock_ints:
block_message: list[str] = []
foriinrange(block_size-1, -1, -1):
iflen(message) +i<message_length:
ascii_number=block_int// (BYTE_SIZE**i)
block_int=block_int% (BYTE_SIZE**i)
block_message.insert(0, chr(ascii_number))
message.extend(block_message)
return"".join(message)


defencrypt_message(
message: str, key: tuple[int, int], block_size: int=DEFAULT_BLOCK_SIZE
) ->list[int]:
encrypted_blocks= []
n, e=key

forblockinget_blocks_from_text(message, block_size):
encrypted_blocks.append(pow(block, e, n))
returnencrypted_blocks


defdecrypt_message(
encrypted_blocks: list[int],
message_length: int,
key: tuple[int, int],
block_size: int=DEFAULT_BLOCK_SIZE,
) ->str:
decrypted_blocks= []
n, d=key
forblockinencrypted_blocks:
decrypted_blocks.append(pow(block, d, n))
returnget_text_from_blocks(decrypted_blocks, message_length, block_size)


defread_key_file(key_filename: str) ->tuple[int, int, int]:
withopen(key_filename) asfo:
content=fo.read()
key_size, n, eor_d=content.split(",")
return (int(key_size), int(n), int(eor_d))


defencrypt_and_write_to_file(
message_filename: str,
key_filename: str,
message: str,
block_size: int=DEFAULT_BLOCK_SIZE,
) ->str:
key_size, n, e=read_key_file(key_filename)
ifkey_size<block_size*8:
sys.exit(
f"ERROR: Block size is {block_size*8} bits and key size is {key_size} "
"bits. The RSA cipher requires the block size to be equal to or greater "
"than the key size. Either decrease the block size or use different keys."
 )

encrypted_blocks= [str(i) foriinencrypt_message(message, (n, e), block_size)]

encrypted_content=",".join(encrypted_blocks)
encrypted_content=f"{len(message)}_{block_size}_{encrypted_content}"
withopen(message_filename, "w") asfo:
fo.write(encrypted_content)
returnencrypted_content


defread_from_file_and_decrypt(message_filename: str, key_filename: str) ->str:
key_size, n, d=read_key_file(key_filename)
withopen(message_filename) asfo:
content=fo.read()
message_length_str, block_size_str, encrypted_message=content.split("_")
message_length=int(message_length_str)
block_size=int(block_size_str)

ifkey_size<block_size*8:
sys.exit(
f"ERROR: Block size is {block_size*8} bits and key size is {key_size} "
"bits. The RSA cipher requires the block size to be equal to or greater "
"than the key size. Were the correct key file and encrypted file specified?"
 )

encrypted_blocks= []
forblockinencrypted_message.split(","):
encrypted_blocks.append(int(block))

returndecrypt_message(encrypted_blocks, message_length, (n, d), block_size)


defmain() ->None:
filename="encrypted_file.txt"
response=input(r"Encrypt\Decrypt [e\d]: ")

ifresponse.lower().startswith("e"):
mode="encrypt"
elifresponse.lower().startswith("d"):
mode="decrypt"

ifmode=="encrypt":
ifnotos.path.exists("rsa_pubkey.txt"):
rkg.make_key_files("rsa", 1024)

message=input("\nEnter message: ")
pubkey_filename="rsa_pubkey.txt"
print(f"Encrypting and writing to {filename}...")
encrypted_text=encrypt_and_write_to_file(filename, pubkey_filename, message)

print("\nEncrypted text:")
print(encrypted_text)

elifmode=="decrypt":
privkey_filename="rsa_privkey.txt"
print(f"Reading from {filename} and decrypting...")
decrypted_text=read_from_file_and_decrypt(filename, privkey_filename)
print("writing decryption to rsa_decryption.txt...")
withopen("rsa_decryption.txt", "w") asdec:
dec.write(decrypted_text)

print("\nDecryption:")
print(decrypted_text)


if__name__=="__main__":
main()