from crypto_lib import *

# Example usage
if __name__ == "__main__":
	generate_new_keys = input('Generate new keys? (y/n): ')
	if generate_new_keys == 'y':
		# 1. Generate keys
		print("1. Generating key pair...\n")
		private_key, public_key = CryptoLib.generate_keys(2048)

		# 2. Save keys to files
		print("2. Saving keys to files...\n")
		CryptoLib.save_private_key(private_key, "private_key.pem")
		CryptoLib.save_public_key(public_key, "public_key.pem")

		encrypted = ''
	else:
		encrypted = CryptoLib.base64_to_bytes(input('Enter message (if exist): '))

	# 6. Load keys from files
	print("3. Loading keys from files...\n")
	loaded_private_key = CryptoLib.load_private_key("private_key.pem")
	loaded_public_key = CryptoLib.load_public_key("public_key.pem")

	# 3. Convert public key to string and back
	print("4. Loaded keys:.")
	private_key_str = CryptoLib.public_key_to_string(loaded_private_key)
	public_key_str = CryptoLib.public_key_to_string(loaded_public_key)
	print(f'Private key as string: {private_key_str[:50]}...')
	print(f'Public key as string: {public_key_str[:50]}...\n')

	print(f'5. Convert public key from string...\n')
	restored_public_key = CryptoLib.public_key_from_string(public_key_str)

	if not encrypted:
		# 6. Encrypt data
		print("6. Encrypting data...")
		message = "Hello, this is a secret message from Python!"
		print(f'Original message: {message}')
		encrypted = CryptoLib.encrypt(loaded_public_key, message)
		print(f'Encrypted (base64): {CryptoLib.bytes_to_base64(encrypted)}')
		print(f"Encrypted data size: {len(encrypted)} bytes\n")

	# 7. Decrypt data
	decrypted_data = CryptoLib.decrypt(loaded_private_key, encrypted)
	decrypted_message = decrypted_data.decode('utf-8')
	print(f"7. Decrypted message: {decrypted_message}")