Kaushik Narayan R c6ecfa305b Initial commit
Completed DH-AES part yesterday
2023-10-29 16:18:39 -07:00

103 lines
3.6 KiB
Python

import os
import argparse
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
def prettyBytes(bytestr: bytes) -> str:
"""Prints byte strings like `\\x3d\\xe9\\xb7` as `3D E9 B7`"""
return " ".join([bytestr[i : i + 1].hex().upper() for i in range(0, len(bytestr))])
def buildExpNum(exponent: int, constant: int, base: int = 2) -> int:
return pow(base, exponent) - constant
def calculateSharedKey(
g_e: int, g_c: int, N_e: int, N_c: int, x: int, gy_modN: int
) -> int:
"""Diffie-Hellman shared key computation (single-party)"""
# g = build_exp_num(g_e, g_c) # not needed since gy_modN is given
N = buildExpNum(N_e, N_c)
return pow(gy_modN, x, N)
def encrypt(plaintext: bytes, key: int, iv: bytearray) -> bytes:
"""Encryption using AES-256 in CBC mode"""
key_byte_length = 32
cipher = Cipher(
algorithms.AES(key.to_bytes(key_byte_length, byteorder="little")), modes.CBC(iv)
)
padder = padding.PKCS7(algorithms.AES.block_size).padder()
padded_plaintext = padder.update(plaintext) + padder.finalize()
encryptor = cipher.encryptor()
return encryptor.update(padded_plaintext) + encryptor.finalize()
def decrypt(ciphertext: bytes, key: int, iv: bytearray) -> bytes:
"""Decryption using AES-256 in CBC mode"""
key_byte_length = 32
cipher = Cipher(
algorithms.AES(key.to_bytes(key_byte_length, byteorder="little")), modes.CBC(iv)
)
decryptor = cipher.decryptor()
return decryptor.update(ciphertext) + decryptor.finalize()
def main(args: argparse.Namespace) -> None:
# pre-processing inputs
iv = bytearray.fromhex(args.initialization_vector)
g_e = int(args.g_e)
g_c = int(args.g_c)
N_e = int(args.N_e)
N_c = int(args.N_c)
x = int(args.x)
gy_modN = int(args.gy_modN)
ciphertext = bytearray.fromhex(args.ciphertext)
plaintext = str.encode(args.plaintext, "utf-8")
verbose = args.verbose
shared_key = calculateSharedKey(g_e, g_c, N_e, N_c, x, gy_modN)
if verbose:
print(("-" * os.get_terminal_size().columns) + "\n")
print("Shared key\t\t", shared_key)
print(("-" * os.get_terminal_size().columns) + "\n")
new_pt = decrypt(ciphertext, shared_key, iv)
if verbose:
print("Given ciphertext\t", prettyBytes(ciphertext))
print("Calculated plaintext\t", new_pt.decode("utf-8"))
print(("-" * os.get_terminal_size().columns) + "\n")
new_ct = encrypt(plaintext, shared_key, iv)
if verbose:
print("Given plaintext\t\t", args.plaintext)
print("Calculated ciphertext\t", prettyBytes(new_ct))
print(("-" * os.get_terminal_size().columns) + "\n")
if not verbose:
print(f"{new_pt.decode('utf-8')}, {prettyBytes(new_ct)}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="DHKE keygen and AES encryption/decryption"
)
parser.add_argument("initialization_vector", help="")
parser.add_argument("g_e", help="Exponent (base 2), for g")
parser.add_argument("g_c", help="Constant to be subtracted, for g")
parser.add_argument("N_e", help="Exponent (base 2), for N")
parser.add_argument("N_c", help="Constant to be subtracted, for N")
parser.add_argument("x", help="Alice's private key value")
parser.add_argument("gy_modN", help="g^y modulo N, computed by Bob")
parser.add_argument("ciphertext", help="Ciphertext to be decrypted")
parser.add_argument("plaintext", help="Plaintext to be encrypted")
parser.add_argument("-v", "--verbose", action="store_true")
args = parser.parse_args()
main(args)