package hsr.dsa.utils;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

public class RSA {

	private static final int RSA_KEY_LENGTH = 4096;
	private static final String ALGORITHM_NAME = "RSA";
	private static final String PADDING_SCHEME = "OAEPWITHSHA-512ANDMGF1PADDING";
	private static final String MODE_OF_OPERATION = "ECB"; // This essentially means none behind the scene

	private KeyPairGenerator rsaKeyGen;
	private KeyPair rsaKeyPair;
	private PublicKey publicKey;
	private PrivateKey privateKey;
	private Cipher cipher;

	public RSA(String pathKeys, String publicKeyName, String privateKeyName)
			throws NoSuchPaddingException, NoSuchAlgorithmException, IOException, InvalidKeySpecException {
		this.cipher = Cipher.getInstance(ALGORITHM_NAME + "/" + MODE_OF_OPERATION + "/" + PADDING_SCHEME);
		if (Files.exists(Paths.get(pathKeys + "_" + publicKeyName))
				&& Files.exists(Paths.get(pathKeys + privateKeyName))) {

			KeyFactory kf = KeyFactory.getInstance("RSA");
			// Read Public Key ###FORMAT X509###
			byte[] publicKeyBytes = Files.readAllBytes(Paths.get(pathKeys + publicKeyName));
			X509EncodedKeySpec specPub = new X509EncodedKeySpec(publicKeyBytes);
			this.publicKey = kf.generatePublic(specPub);

			// Read Private Key ###FORMAT PKCS8###
			byte[] privateKeyBytes = Files.readAllBytes(Paths.get(pathKeys + privateKeyName));
			PKCS8EncodedKeySpec specP = new PKCS8EncodedKeySpec(privateKeyBytes);
			this.privateKey = kf.generatePrivate(specP);
		} else {
			// Generate Key Pairs
			this.rsaKeyGen = KeyPairGenerator.getInstance(ALGORITHM_NAME);
			this.rsaKeyGen.initialize(RSA_KEY_LENGTH);
			this.rsaKeyPair = this.rsaKeyGen.generateKeyPair();
			this.publicKey = this.rsaKeyPair.getPublic();
			this.privateKey = this.rsaKeyPair.getPrivate();

			// Write keys to files
			Files.write(Paths.get(pathKeys + publicKeyName), this.rsaKeyPair.getPublic().getEncoded());
			Files.write(Paths.get(pathKeys + privateKeyName), this.rsaKeyPair.getPrivate().getEncoded());
		}
	}

	public RSA(PublicKey publicKey) throws NoSuchAlgorithmException, NoSuchPaddingException {
		this.cipher = Cipher.getInstance(ALGORITHM_NAME + "/" + MODE_OF_OPERATION + "/" + PADDING_SCHEME);
		this.publicKey = publicKey;
	}

	public RSA() throws NoSuchAlgorithmException, NoSuchPaddingException {
		this.cipher = Cipher.getInstance(ALGORITHM_NAME + "/" + MODE_OF_OPERATION + "/" + PADDING_SCHEME);
	}

	public PublicKey getPublicKey() {
		return this.publicKey;
	}

	public void setPublicKey(PublicKey publicKey) {
		this.publicKey = publicKey;
	}

	public byte[] encrypt(byte[] message) throws BadPaddingException, IllegalBlockSizeException, InvalidKeyException {
		this.cipher.init(Cipher.ENCRYPT_MODE, this.publicKey);
		byte[] cipherTextArray = this.cipher.doFinal(message);
		return cipherTextArray;
	}

	public String encrypt(String message) throws BadPaddingException, IllegalBlockSizeException, InvalidKeyException {
		return Base64.getEncoder().encodeToString(this.encrypt(message.getBytes()));
	}

	public byte[] decrypt(byte[] encryptedMessage)
			throws InvalidKeyException, BadPaddingException, IllegalBlockSizeException {
		this.cipher.init(Cipher.DECRYPT_MODE, this.privateKey);
		return this.cipher.doFinal(encryptedMessage);
	}

	public String decrypt(String encryptedMessage)
			throws BadPaddingException, InvalidKeyException, IllegalBlockSizeException {
		return new String(this.decrypt(Base64.getDecoder().decode(encryptedMessage)));
	}
}
