package launcher.helper;

import launcher.LauncherAPI;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.nio.file.Path;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.CertificateEncodingException;
import java.security.interfaces.RSAKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.jar.JarFile;

public final class SecurityHelper
{
    // Algorithm constants
    @LauncherAPI
    public static final String RSA_ALGO = "RSA";
    @LauncherAPI
    public static final String RSA_SIGN_ALGO = "SHA256withRSA";
    @LauncherAPI
    public static final String RSA_CIPHER_ALGO = "RSA/ECB/PKCS1Padding";

    // Algorithm size constants
    @LauncherAPI
    public static final int TOKEN_LENGTH = 16;
    @LauncherAPI
    public static final int TOKEN_STRING_LENGTH = TOKEN_LENGTH << 1;
    @LauncherAPI
    public static final int RSA_KEY_LENGTH_BITS = 2048;
    @LauncherAPI
    public static final int RSA_KEY_LENGTH = RSA_KEY_LENGTH_BITS / Byte.SIZE;
    @LauncherAPI
    public static final int CRYPTO_MAX_LENGTH = 2048;

    // Certificate constants
    @LauncherAPI
    public static final String CERTIFICATE_DIGEST = "b87c079e3bf6e709860e05e283678c857b6a27916c2ba280a212f78f1a2ec20a";
    @LauncherAPI
    public static final String HEX = "0123456789abcdef";

    // Random generator constants
    private static final char[] VOWELS = {'e', 'u', 'i', 'o', 'a'};
    private static final char[] CONS = {'r', 't', 'p', 's', 'd', 'f', 'g', 'h', 'k', 'l', 'c', 'v', 'b', 'n', 'm'};

    private SecurityHelper()
    {
    }

    @LauncherAPI
    public static byte[] digest(DigestAlgorithm algo, String s)
    {
        return digest(algo, IOHelper.encode(s));
    }

    @LauncherAPI
    public static byte[] digest(DigestAlgorithm algo, URL url) throws IOException
    {
        try (InputStream input = IOHelper.newInput(url))
        {
            return digest(algo, input);
        }
    }

    @LauncherAPI
    public static byte[] digest(DigestAlgorithm algo, Path file) throws IOException
    {
        try (InputStream input = IOHelper.newInput(file))
        {
            return digest(algo, input);
        }
    }

    @LauncherAPI
    public static byte[] digest(DigestAlgorithm algo, byte[] bytes)
    {
        return newDigest(algo).digest(bytes);
    }

    @LauncherAPI
    public static byte[] digest(DigestAlgorithm algo, InputStream input) throws IOException
    {
        byte[] buffer = IOHelper.newBuffer();
        MessageDigest digest = newDigest(algo);
        for (int length = input.read(buffer); length != -1; length = input.read(buffer))
        {
            digest.update(buffer, 0, length);
        }
        return digest.digest();
    }

    @LauncherAPI
    public static KeyPair genRSAKeyPair(SecureRandom random)
    {
        try
        {
            KeyPairGenerator generator = KeyPairGenerator.getInstance(RSA_ALGO);
            generator.initialize(RSA_KEY_LENGTH_BITS, random);
            return generator.genKeyPair();
        }
        catch (NoSuchAlgorithmException e)
        {
            throw new InternalError(e);
        }
    }

    @LauncherAPI
    public static KeyPair genRSAKeyPair()
    {
        return genRSAKeyPair(newRandom());
    }

    @LauncherAPI
    public static boolean isValidCertificate(Certificate cert)
    {
        try
        {
            return toHex(digest(DigestAlgorithm.SHA256, cert.getEncoded())).equals(CERTIFICATE_DIGEST);
        }
        catch (CertificateEncodingException e)
        {
            throw new InternalError(e);
        }
    }

    @LauncherAPI
    public static boolean isValidCertificates(Certificate... certs)
    {
        return certs != null && certs.length == 1 && isValidCertificate(certs[0]);
    }

    @LauncherAPI
    public static boolean isValidCertificates(Class<?> clazz)
    {
        // Verify META-INF/MANIFEST.MF certificate
        Certificate[] certificates = JVMHelper.getCertificates(JarFile.MANIFEST_NAME);
        if (certificates == null || !isValidCertificates(certificates))
        {
            return false;
        }

        // Verify class certificate
        CodeSource source = clazz.getProtectionDomain().getCodeSource();
        return source != null && isValidCertificates(source.getCertificates());
    }

    @LauncherAPI
    public static boolean isValidSign(Path path, byte[] sign, RSAPublicKey publicKey) throws IOException, SignatureException
    {
        try (InputStream input = IOHelper.newInput(path))
        {
            return isValidSign(input, sign, publicKey);
        }
    }

    @LauncherAPI
    public static boolean isValidSign(byte[] bytes, byte[] sign, RSAPublicKey publicKey) throws SignatureException
    {
        Signature signature = newRSAVerifySignature(publicKey);
        try
        {
            signature.update(bytes);
        }
        catch (SignatureException e)
        {
            throw new InternalError(e);
        }
        return signature.verify(sign);
    }

    @LauncherAPI
    public static boolean isValidSign(InputStream input, byte[] sign, RSAPublicKey publicKey) throws IOException, SignatureException
    {
        Signature signature = newRSAVerifySignature(publicKey);
        updateSignature(input, signature);
        return signature.verify(sign);
    }

    @LauncherAPI
    public static boolean isValidSign(URL url, byte[] sign, RSAPublicKey publicKey) throws IOException, SignatureException
    {
        try (InputStream input = IOHelper.newInput(url))
        {
            return isValidSign(input, sign, publicKey);
        }
    }

    @LauncherAPI
    public static boolean isValidToken(CharSequence token)
    {
        return token.length() == TOKEN_STRING_LENGTH && token.chars().allMatch(ch -> HEX.indexOf(ch) >= 0);
    }

    @LauncherAPI
    public static MessageDigest newDigest(DigestAlgorithm algo)
    {
        VerifyHelper.verify(algo, a -> a != DigestAlgorithm.PLAIN, "PLAIN digest");
        try
        {
            return MessageDigest.getInstance(algo.name);
        }
        catch (NoSuchAlgorithmException e)
        {
            throw new InternalError(e);
        }
    }

    @LauncherAPI
    public static Cipher newRSADecryptCipher(RSAPrivateKey key)
    {
        return newRSACipher(Cipher.DECRYPT_MODE, key);
    }

    @LauncherAPI
    public static Cipher newRSAEncryptCipher(RSAPublicKey key)
    {
        return newRSACipher(Cipher.ENCRYPT_MODE, key);
    }

    @LauncherAPI
    public static Signature newRSASignSignature(RSAPrivateKey key)
    {
        Signature signature = newRSASignature();
        try
        {
            signature.initSign(key);
        }
        catch (InvalidKeyException e)
        {
            throw new InternalError(e);
        }
        return signature;
    }

    @LauncherAPI
    public static Signature newRSAVerifySignature(RSAPublicKey key)
    {
        Signature signature = newRSASignature();
        try
        {
            signature.initVerify(key);
        }
        catch (InvalidKeyException e)
        {
            throw new InternalError(e);
        }
        return signature;
    }

    @LauncherAPI
    public static SecureRandom newRandom()
    {
        return new SecureRandom();
    }

    @LauncherAPI
    public static byte[] randomBytes(Random random, int length)
    {
        byte[] bytes = new byte[length];
        random.nextBytes(bytes);
        return bytes;
    }

    @LauncherAPI
    public static byte[] randomBytes(int length)
    {
        return randomBytes(newRandom(), length);
    }

    @LauncherAPI
    public static String randomStringToken(Random random)
    {
        return toHex(randomToken(random));
    }

    @LauncherAPI
    public static String randomStringToken()
    {
        return randomStringToken(newRandom());
    }

    @LauncherAPI
    public static byte[] randomToken(Random random)
    {
        return randomBytes(random, TOKEN_LENGTH);
    }

    @LauncherAPI
    public static byte[] randomToken()
    {
        return randomToken(newRandom());
    }

    @LauncherAPI
    public static String randomUsername(Random random)
    {
        int usernameLength = 3 + random.nextInt(7); // 3-9

        // Choose prefix
        String prefix;
        int prefixType = random.nextInt(7);
        if (usernameLength >= 5 && prefixType == 6)
        { // (6) 2-char
            prefix = random.nextBoolean() ? "Mr" : "Dr";
            usernameLength -= 2;
        }
        else if (usernameLength >= 6 && prefixType == 5)
        { // (5) 3-char
            prefix = "Mrs";
            usernameLength -= 3;
        }
        else
        {
            prefix = "";
        }

        // Choose suffix
        String suffix;
        int suffixType = random.nextInt(7); // 0-6, 7 values
        if (usernameLength >= 5 && suffixType == 6)
        { // (6) 10-99
            suffix = String.valueOf(10 + random.nextInt(90));
            usernameLength -= 2;
        }
        else if (usernameLength >= 7 && suffixType == 5)
        { // (5) 1990-2015
            suffix = String.valueOf(1990 + random.nextInt(26));
            usernameLength -= 4;
        }
        else
        {
            suffix = "";
        }

        // Choose name
        int consRepeat = 0;
        boolean consPrev = random.nextBoolean();
        char[] chars = new char[usernameLength];
        for (int i = 0; i < chars.length; i++)
        {
            if (i > 1 && consPrev && random.nextInt(10) == 0)
            { // Doubled
                chars[i] = chars[i - 1];
                continue;
            }

            // Choose next char
            if (consRepeat < 1 && random.nextInt() == 5)
            {
                consRepeat++;
            }
            else
            {
                consRepeat = 0;
                consPrev ^= true;
            }

            // Choose char
            char[] alphabet = consPrev ? CONS : VOWELS;
            chars[i] = alphabet[random.nextInt(alphabet.length)];
        }

        // Make first letter uppercase
        if (!prefix.isEmpty() || random.nextBoolean())
        {
            chars[0] = Character.toUpperCase(chars[0]);
        }

        // Return chosen name (and verify for sure)
        return VerifyHelper.verifyUsername(prefix + new String(chars) + suffix);
    }

    @LauncherAPI
    public static String randomUsername()
    {
        return randomUsername(newRandom());
    }

    @LauncherAPI
    public static byte[] sign(InputStream input, RSAPrivateKey privateKey) throws IOException
    {
        Signature signature = newRSASignSignature(privateKey);
        updateSignature(input, signature);
        try
        {
            return signature.sign();
        }
        catch (SignatureException e)
        {
            throw new InternalError(e);
        }
    }

    @LauncherAPI
    public static byte[] sign(byte[] bytes, RSAPrivateKey privateKey)
    {
        Signature signature = newRSASignSignature(privateKey);
        try
        {
            signature.update(bytes);
            return signature.sign();
        }
        catch (SignatureException e)
        {
            throw new InternalError(e);
        }
    }

    @LauncherAPI
    public static byte[] sign(Path path, RSAPrivateKey privateKey) throws IOException
    {
        try (InputStream input = IOHelper.newInput(path))
        {
            return sign(input, privateKey);
        }
    }

    @LauncherAPI
    public static String toHex(byte[] bytes)
    {
        int offset = 0;
        char[] hex = new char[bytes.length << 1];
        for (byte currentByte : bytes)
        {
            int ub = Byte.toUnsignedInt(currentByte);
            hex[offset] = HEX.charAt(ub >>> 4);
            offset++;
            hex[offset] = HEX.charAt(ub & 0x0F);
            offset++;
        }
        return new String(hex);
    }

    @LauncherAPI
    public static RSAPrivateKey toPrivateRSAKey(byte[] bytes) throws InvalidKeySpecException
    {
        return (RSAPrivateKey) newRSAKeyFactory().generatePrivate(new PKCS8EncodedKeySpec(bytes));
    }

    @LauncherAPI
    public static RSAPublicKey toPublicRSAKey(byte[] bytes) throws InvalidKeySpecException
    {
        return (RSAPublicKey) newRSAKeyFactory().generatePublic(new X509EncodedKeySpec(bytes));
    }

    @LauncherAPI
    public static void verifyCertificates(Class<?> clazz)
    {
        if (!isValidCertificates(clazz))
        {
            throw new SecurityException("Invalid certificates");
        }
    }

    @LauncherAPI
    public static void verifySign(byte[] bytes, byte[] sign, RSAPublicKey publicKey) throws SignatureException
    {
        if (!isValidSign(bytes, sign, publicKey))
        {
            throw new SignatureException("Invalid sign");
        }
    }

    @LauncherAPI
    public static void verifySign(InputStream input, byte[] sign, RSAPublicKey publicKey) throws SignatureException, IOException
    {
        if (!isValidSign(input, sign, publicKey))
        {
            throw new SignatureException("Invalid stream sign");
        }
    }

    @LauncherAPI
    public static void verifySign(Path path, byte[] sign, RSAPublicKey publicKey) throws SignatureException, IOException
    {
        if (!isValidSign(path, sign, publicKey))
        {
            throw new SignatureException(String.format("Invalid file sign: '%s'", path));
        }
    }

    @LauncherAPI
    public static void verifySign(URL url, byte[] sign, RSAPublicKey publicKey) throws SignatureException, IOException
    {
        if (!isValidSign(url, sign, publicKey))
        {
            throw new SignatureException(String.format("Invalid URL sign: '%s'", url));
        }
    }

    @LauncherAPI
    public static String verifyToken(String token)
    {
        return VerifyHelper.verify(token, SecurityHelper::isValidToken, String.format("Invalid token: '%s'", token));
    }

    private static Cipher newCipher(String algo)
    {
        // IDK Why, but collapsing catch blocks makes ProGuard generate invalid stackmap
        try
        {
            return Cipher.getInstance(algo);
        }
        catch (NoSuchAlgorithmException e)
        {
            throw new InternalError(e);
        }
        catch (NoSuchPaddingException e)
        {
            throw new InternalError(e);
        }
    }

    private static Cipher newRSACipher(int mode, RSAKey key)
    {
        Cipher cipher = newCipher(RSA_CIPHER_ALGO);
        try
        {
            cipher.init(mode, (Key) key);
        }
        catch (InvalidKeyException e)
        {
            throw new InternalError(e);
        }
        return cipher;
    }

    private static KeyFactory newRSAKeyFactory()
    {
        try
        {
            return KeyFactory.getInstance(RSA_ALGO);
        }
        catch (NoSuchAlgorithmException e)
        {
            throw new InternalError(e);
        }
    }

    private static Signature newRSASignature()
    {
        try
        {
            return Signature.getInstance(RSA_SIGN_ALGO);
        }
        catch (NoSuchAlgorithmException e)
        {
            throw new InternalError(e);
        }
    }

    private static void updateSignature(InputStream input, Signature signature) throws IOException
    {
        byte[] buffer = IOHelper.newBuffer();
        for (int length = input.read(buffer); length >= 0; length = input.read(buffer))
        {
            try
            {
                signature.update(buffer, 0, length);
            }
            catch (SignatureException e)
            {
                throw new InternalError(e);
            }
        }
    }

    @LauncherAPI
    public enum DigestAlgorithm
    {
        PLAIN("plain", -1), MD5("MD5", 128), SHA1("SHA-1", 160), SHA224("SHA-224", 224), SHA256("SHA-256", 256), SHA512("SHA-512", 512);
        private static final Map<String, DigestAlgorithm> ALGORITHMS;

        static
        {
            DigestAlgorithm[] algorithmsValues = values();
            ALGORITHMS = new HashMap<>(algorithmsValues.length);
            for (DigestAlgorithm algorithm : algorithmsValues)
            {
                ALGORITHMS.put(algorithm.name, algorithm);
            }
        }

        // Instance
        public final String name;
        public final int bits;
        public final int bytes;

        DigestAlgorithm(String name, int bits)
        {
            this.name = name;
            this.bits = bits;

            // Convert to bytes
            bytes = bits / Byte.SIZE;
            assert bits % Byte.SIZE == 0;
        }

        public static DigestAlgorithm byName(String name)
        {
            return VerifyHelper.getMapValue(ALGORITHMS, name, String.format("Unknown digest algorithm: '%s'", name));
        }

        @Override
        public String toString()
        {
            return name;
        }

        public byte[] verify(byte[] digest)
        {
            if (digest.length != bytes)
            {
                throw new IllegalArgumentException("Invalid digest length: " + digest.length);
            }
            return digest;
        }
    }
}