package com.example.crypto;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;

public class CryptoUtil {

    public static String encode(String data, String secretKey) throws Exception {
        Key key = getKey(secretKey.getBytes());
        Cipher cipher = Cipher.getInstance("AES/ECB/PKCS5Padding");
        cipher.init(Cipher.ENCRYPT_MODE, key);
        byte[] bytes = cipher.doFinal(data.getBytes("UTF-8"));
        return Base64.getEncoder().encodeToString(bytes);
    }

    public static String decrypt(String data, String secretKey) throws Exception {
        byte[] decode = Base64.getDecoder().decode(data);
        Key key = getKey(secretKey.getBytes());
        Cipher cipher = Cipher.getInstance("AES/ECB/PKCS5Padding");
        cipher.init(Cipher.DECRYPT_MODE, key);
        byte[] bytes = cipher.doFinal(decode);
        return new String(bytes, "UTF-8");
    }

    private static Key getKey(byte[] secret) throws NoSuchAlgorithmException {
        KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
        SecureRandom secureRandom = SecureRandom.getInstance("SHA1PRNG");
        secureRandom.setSeed(secret);
        keyGenerator.init(128, secureRandom);
        SecretKey secretKey = keyGenerator.generateKey();
        return new SecretKeySpec(secretKey.getEncoded(), "AES");
    }

    public static void main(String[] args) throws Exception {
        if (args.length < 3) {
            System.out.println("Usage: java CryptoUtil <encode/decode> <message> <secretKey>");
            return;
        }

        String operation = args[0];
        String message = args[1];
        String secretKey = args[2];

        if ("encode".equalsIgnoreCase(operation)) {
            System.out.println(encode(message, secretKey));
        } else if ("decode".equalsIgnoreCase(operation)) {
            System.out.println(decrypt(message, secretKey));
        } else {
            System.out.println("Unknown operation: " + operation);
        }
    }
}