import botan2 as botan
import base64

class CryptoLib:
	@staticmethod
	def generate_keys(key_size=2048):
		"""Generate RSA key pair"""
		rng = botan.RandomNumberGenerator()
		private_key = botan.PrivateKey.create("RSA", str(key_size), rng)
		public_key = private_key.get_public_key()
		return private_key, public_key

	@staticmethod
	def encrypt(public_key, data):
		"""Encrypt data with public key"""
		if isinstance(data, str):
			data = data.encode('utf-8')
		rng = botan.RandomNumberGenerator()
		encryptor = botan.PKEncrypt(public_key, "OAEP(SHA-256)")
		return encryptor.encrypt(data, rng)

	@staticmethod
	def decrypt(private_key, encrypted_data):
		"""Decrypt data with private key"""
		decryptor = botan.PKDecrypt(private_key, "OAEP(SHA-256)")
		return decryptor.decrypt(encrypted_data)

	@staticmethod
	def public_key_to_string(public_key):
		"""Convert public key to PEM string"""
		return public_key.to_pem().encode('utf-8')

	@staticmethod
	def public_key_from_string(key_string):
		"""Load public key from PEM string"""
		if isinstance(key_string, str):
			key_string = key_string.encode('utf-8')
		return botan.PublicKey.load(key_string)

	@staticmethod
	def private_key_to_string(private_key):
		"""Convert private key to PEM string"""
		return private_key.to_pem().decode('utf-8')

	@staticmethod
	def private_key_from_string(key_string):
		"""Load private key from PEM string"""
		if isinstance(key_string, str):
			key_string = key_string.encode('utf-8')
		return botan.PrivateKey.load(key_string)

	@staticmethod
	def save_private_key(private_key, filename):
		"""Save private key to file"""
		with open(filename, 'w') as f:
			# f.write(private_key.to_pem().encode('utf-8'))
			f.write(private_key.to_pem())

	@staticmethod
	def save_public_key(public_key, filename):
		"""Save public key to file"""
		with open(filename, 'w') as f:
			# f.write(public_key.to_pem().decode('utf-8'))
			f.write(public_key.to_pem())

	@staticmethod
	def load_private_key(filename):
		"""Load private key from file"""
		with open(filename, 'r') as f:
			pem_data = f.read()
		return botan.PrivateKey.load(pem_data.encode('utf-8'))

	@staticmethod
	def load_public_key(filename):
		"""Load public key from file"""
		with open(filename, 'r') as f:
			pem_data = f.read()
		return botan.PublicKey.load(pem_data.encode('utf-8'))

	@staticmethod
	def bytes_to_base64(public_key: bytes) -> str:
		return base64.b64encode(public_key).decode('utf-8')

	@staticmethod
	def base64_to_bytes(key_string: str) -> bytes:
		return base64.b64decode(key_string.encode('utf-8'))