def process(self, request: bytes) -> bool:
"""Processes a request. Returns a boolean representing whether the
message was valid and successfully processed. Each request looks like:
message || IV || MAC
"""
assert len(
request
) > BLOCK_SIZE * 2, 'Request must be at least two blocks long'
MSG = request[:-BLOCK_SIZE * 2]
IV = request[-BLOCK_SIZE * 2:-BLOCK_SIZE]
MAC = request[-BLOCK_SIZE:]
expected_mac = cbc_encrypt(MSG, key=self.K, iv=IV)[-BLOCK_SIZE:]
if MAC != expected_mac:
return False
fields = self.get_fields(MSG)
amount = int(fields['amount'])
sender = int(fields['from'])
receiver = int(fields['to'])
# NOTE: I don't do _any_ checking to see whether or not an account has
# funds. Who cares about realism?
self.accounts[sender] -= amount
self.accounts[receiver] += amount
return True
def randcrypt():
"""
Selects one of the strings above, at random, encrypts it using our random
global key, and returns the ciphertext and iv.
"""
msg = choice(POSSIBLE_INPUTS)
iv = os.urandom(16)
ciphertext = cbc_encrypt(msg, KEY, iv)
return ciphertext, iv
def recieve(self, msg: bytes, iv: bytes) -> bytes:
"""Recieves a message, decrypts it with A's key, and then re-encrypts
with a new IV and returns
Returns:
(bytes, bytes): msg, iv tuple
"""
s = modpow(self.A, self.b, self.p)
k = dh_digest(s)
plaintext = pkcs7_strip(cbc_decrypt(msg, k, iv))
iv = os.urandom(16)
return cbc_encrypt(plaintext, k, iv), iv
def encrypt(msg: bytes, key: bytes = None, iv: bytes = None) -> bytes:
""" Sandwiches user input between two existing strings, surrounding any ';'
or '=' characters in single quotes. Pads and encrypts in AES CBC using the
random key and a random IV, then returns the encrypted messge.
"""
# Fall back to module defaults for local tests
key = key or KEY
iv = iv or IV
# I interpreted "quote out" to mean surround disallowed characters with
# single quotes
msg_str = msg.decode()
msg = CLEAN_RE.sub(r"'\1'", msg_str)
# Convert string (necessary for regex replacement) back to bytes
msg = bytes(msg.encode())
# Tack on the prefix and suffix
msg = PREFIX + msg + SUFFIX
return cbc_encrypt(msg, key, iv)
def randcryptor(msg: bytes) -> bytes:
"""
Encrypts the input with AES in either ECB or CBC, randomly.
"""
# Pad the message with 5-10 random bytes
padding = os.urandom(randint(5, 10))
padded_msg = pkcs7_pad(padding + msg + padding, BLOCK_SIZE)
# Generate a random key
key = generate_aes_key()
# Randomly choose whether to use ECB or CBC
mode = 'cbc' if bool(randint(0, 1)) else 'ecb'
if mode == 'cbc':
iv = os.urandom(BLOCK_SIZE)
encrypted_msg = cbc_encrypt(padded_msg, key, iv)
else: # ECB
cipher = AES.new(key, AES.MODE_ECB)
encrypted_msg = cipher.encrypt(padded_msg)
return encrypted_msg, mode # FOR TESTING ONLY
print(
'Challenge #34 - Implement a MITM key-fixing attack on Diffie-Hellman with parameter injection'
)
p = 12345
g = 5
A, a = dh_key(p, g)
echo = Echo(p, g, A)
# Test basic echo
plaintext = b'some plaintext'
key = dh_digest(modpow(echo.B, a, p))
iv = os.urandom(16)
ciphertext = cbc_encrypt(plaintext, key, iv)
resp_ciphertext, resp_iv = echo.recieve(ciphertext, iv)
# The sent and recieved plaintexts should be identical, after padding is
# stripped away
resp_plaintext = pkcs7_strip(cbc_decrypt(resp_ciphertext, key, resp_iv))
assert plaintext == resp_plaintext
# MITM
mitm = Middleman(p, g, A)
plaintext = b'top secret, do not read'
key = dh_digest(modpow(mitm.B, a, p))
iv = os.urandom(16)
ciphertext = cbc_encrypt(plaintext, key, iv)
resp_ciphertext, resp_iv = mitm.recieve(ciphertext, iv)
def generate_mac(self, msg: bytes) -> bytes:
"""Returns the MAC of a given message.
"""
return cbc_encrypt(msg, key=self.K, iv=self.IV)[-BLOCK_SIZE:]