Added test

Added a test for the encryption mode

Author: q3st1on

tests/test_padding_oracle.py

@@ -1,3 +1,4 @@
+from cryptography.hazmat.primitives import padding
 from padding_oracle import padding_oracle
 from .cryptor import VulnerableCryptor
 
@@ -14,6 +15,21 @@ def test_padding_oracle_basic():
 
     assert decrypted == plaintext
 
+def test_padding_oracle_encryption():
+    cryptor = VulnerableCryptor()
+
+    plaintext = b'the quick brown fox jumps over the lazy dog'
+    ciphertext = cryptor.encrypt(plaintext)
+
+    padder = padding.PKCS7(128).padder()
+    payload = padder.update(plaintext) + padder.finalize()
+
+    encrypted = padding_oracle(payload, cryptor.block_size,
+                               cryptor.oracle, 4, null_byte=b'?', mode='encrypt')
+    print(encrypted)
+
+    assert encrypted == ciphertext
 
 if __name__ == '__main__':
     test_padding_oracle_basic()
+    test_padding_oracle_encryption()