#-*- coding: utf-8 -*-
# (c) 2015 Anders Andersen
# See http://www.cs.uit.no/~aa/dist/tools/py/COPYING for details


# Shared values (both examples encrypt/decrypt)
pwd = "My secret password"
txt = "Attack at dawn"


# The ../pycrypto/pycryptex.py example with the cryptography module (see
# README) using the Fernet class (a high level abstraction)

# Load modules
from cryptography.fernet import Fernet
print("\nVersion 1:")

# Generate key and encrypt
key = Fernet.generate_key()
cip = Fernet(key)
msg = cip.encrypt(txt.encode())

# Print text and cipher text
print(txt)
print(msg)

# Decrypt (NB: same copy of key)
cip2 = Fernet(key)
print(cip2.decrypt(msg).decode())


# The example above does not match the "../pycrypto/pycryptex.py" example
# perfectly (no password and different mode and key size). Below is the same
# example without the Fernet class and equal to the PyCrypto version in
# "../pycrypto/pycryptex.py".

# Load modules
import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
print("\nVersion 2:")

# Shared values (encrypt/decrypt)
icv = os.urandom(algorithms.AES.block_size//8)

# Generate key and encrypt
dig = hashes.Hash(hashes.SHA256(), backend=default_backend())
dig.update(pwd.encode())
key = dig.finalize()
cip = Cipher(algorithms.AES(key), modes.CTR(icv), backend=default_backend())
enc = cip.encryptor()
msg = enc.update(txt.encode()) + enc.finalize()

# Print text and cipher text
print(txt)
print(msg)

# Generate key and decrypt
dig2 = hashes.Hash(hashes.SHA256(), backend=default_backend())
dig2.update(pwd.encode())
key2 = dig2.finalize()
cip2 = Cipher(algorithms.AES(key2), modes.CTR(icv), backend=default_backend())
enc2 = cip2.decryptor()
print((enc2.update(msg) + enc2.finalize()).decode() + "\n")