import os

from decryption_functs import *
from encryption_functs import *
from key_pair import *

def test_encryption_no_pwd():
	# create a file to encrypt
	with open("test.txt", "w") as f:
		f.write("Hello, World!")

	# generate a key pair
	generate_key_pair('public.pem', 'private.pem', '2048')

	# load the public and private keys
	public_key = load_public_key("public.pem")
	private_key = load_private_key("private.pem")

	# encrypt the file
	encrypt_file(public_key, "test.txt", "test.enc")

	# decrypt the file
	decrypt_file(private_key, "test.enc", "test.dec")

	# check that the decrypted file is the same as the original
	with open("test.dec", "r") as f:
		assert f.read() == "Hello, World!"

	# cleanup
	os.remove("test.txt")
	os.remove("test.enc")
	os.remove("test.dec")
	os.remove("public.pem")
	os.remove("private.pem")


def test_encryption_with_pwd():
	# create a file to encrypt
	with open("test.txt", "w") as f:
		f.write("Hello, World!")

	# generate a key pair
	generate_key_pair('public.pem', 'private.pem', '2048', 'password')

	# load the public and private keys
	public_key = load_public_key("public.pem")
	private_key = load_private_key("private.pem", 'password')

	# encrypt the file
	encrypt_file(public_key, "test.txt", "test.enc")

	# decrypt the file
	decrypt_file(private_key, "test.enc", "test.dec")

	# check that the decrypted file is the same as the original
	with open("test.dec", "r") as f:
		assert f.read() == "Hello, World!"

	# remove the files
	os.remove("test.txt")
	os.remove("test.enc")
	os.remove("test.dec")
	os.remove("public.pem")
	os.remove("private.pem")


def test_load_private_key_wrong_pwd():
	# generate a key pair
	generate_key_pair('public.pem', 'private.pem', '2048', 'password')

	# try to load the private key with the wrong password
	try:
		load_private_key("private.pem", 'wrong_password')
	except ValueError as e:
		assert str(e) == "Error: The password is not valid."


def test_1mb_file_with_pwd():
	# create a 1mb file to encrypt
	os.system("dd if=/dev/urandom of=test.txt bs=1024 count=1000  >/dev/null 2>&1")

	# generate a key pair
	generate_key_pair('public.pem', 'private.pem', '2048', 'password')

	# load the public and private keys
	public_key = load_public_key("public.pem")
	private_key = load_private_key("private.pem", 'password')

	# encrypt the file
	encrypt_file(public_key, "test.txt", "test.enc")

	# decrypt the file
	decrypt_file(private_key, "test.enc", "test.dec")

	# check that the decrypted file is the same as the original
	assert open("test.txt", "rb").read() == open("test.dec", "rb").read()

	# remove the files
	os.remove("test.txt")
	os.remove("test.enc")
	os.remove("test.dec")
	os.remove("public.pem")
	os.remove("private.pem")


def test_1mb_file_no_pwd():
	# create a 1mb file to encrypt
	os.system("dd if=/dev/urandom of=test.txt bs=1024 count=1000  >/dev/null 2>&1")

	# generate a key pair
	generate_key_pair('public.pem', 'private.pem', '2048')

	# load the public and private keys
	public_key = load_public_key("public.pem")
	private_key = load_private_key("private.pem")

	# encrypt the file
	encrypt_file(public_key, "test.txt", "test.enc")

	# decrypt the file
	decrypt_file(private_key, "test.enc", "test.dec")

	# check that the decrypted file is the same as the original
	assert open("test.txt", "rb").read() == open("test.dec", "rb").read()

	# remove the files
	os.remove("test.txt")
	os.remove("test.enc")
	os.remove("test.dec")
	os.remove("public.pem")
	os.remove("private.pem")


def test_100mb_file_with_pwd():
	# create a 100mb file to encrypt
	os.system("dd if=/dev/urandom of=test.txt bs=1024 count=100000  >/dev/null 2>&1")

	# generate a key pair
	generate_key_pair('public.pem', 'private.pem', '2048', 'password')

	# load the public and private keys
	public_key = load_public_key("public.pem")
	private_key = load_private_key("private.pem", 'password')

	# encrypt the file
	encrypt_file(public_key, "test.txt", "test.enc")

	# decrypt the file
	decrypt_file(private_key, "test.enc", "test.dec")

	# check that the decrypted file is the same as the original
	assert open("test.txt", "rb").read() == open("test.dec", "rb").read()

	# remove the files
	os.remove("test.txt")
	os.remove("test.enc")
	os.remove("test.dec")
	os.remove("public.pem")
	os.remove("private.pem")


def test_100mb_file_no_pwd():
	# create a 100mb file to encrypt
	os.system("dd if=/dev/urandom of=test.txt bs=1024 count=100000  >/dev/null 2>&1")

	# generate a key pair
	generate_key_pair('public.pem', 'private.pem', '2048')

	# load the public and private keys
	public_key = load_public_key("public.pem")
	private_key = load_private_key("private.pem")

	# encrypt the file
	encrypt_file(public_key, "test.txt", "test.enc")

	# decrypt the file
	decrypt_file(private_key, "test.enc", "test.dec")

	# check that the decrypted file is the same as the original
	assert open("test.txt", "rb").read() == open("test.dec", "rb").read()

	# remove the files
	os.remove("test.txt")
	os.remove("test.enc")
	os.remove("test.dec")
	os.remove("public.pem")
	os.remove("private.pem")