package com.atlassian.bitbucket.internal.mirroring.ssh.encoding;

import org.apache.commons.lang.StringUtils;
import org.bouncycastle.util.encoders.Base64;

import javax.annotation.Nonnull;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.security.KeyFactory;
import java.security.PublicKey;
import java.security.interfaces.DSAPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.X509EncodedKeySpec;

import static com.google.common.base.Preconditions.checkNotNull;

/**
 * Encodes/decodes a public key in various formats
 */
public class DefaultPublicKeyEncodingHelper implements PublicKeyEncodingHelper {

    private static final String DSA_TYPE = "ssh-dss";
    private static final String DSA_PREFIX = DSA_TYPE + " ";

    private static final String RSA_TYPE = "ssh-rsa";
    private static final String RSA_PREFIX = RSA_TYPE + " ";

    @Nonnull
    @Override
    public PublicKey decodeBytesFromBase64(@Nonnull String keyText, @Nonnull String algorithm) {
        try {
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64.decode(keyText));
            KeyFactory keyFactory = KeyFactory.getInstance(algorithm);
            return keyFactory.generatePublic(keySpec);
        } catch (Exception e) {
            throw new KeyCodecException(e);
        }
    }

    @Nonnull
    @Override
    public String encodeBytesAsBase64(@Nonnull PublicKey publicKey) {
        return Base64.toBase64String(publicKey.getEncoded());
    }

    @Nonnull
    @Override
    public String encodeAsOpenSsh(@Nonnull PublicKey key) {
        return encodeAsOpenSsh(key, "");
    }

    @Nonnull
    @Override
    public String encodeAsOpenSsh(@Nonnull PublicKey key, @Nonnull String label) {
        checkNotNull(key, "key");
        checkNotNull(label, "label");

        String openSshKey;
        switch (key.getAlgorithm()) {
            case "RSA":
                openSshKey = RSA_PREFIX + Base64.toBase64String(encodeRSA((RSAPublicKey) key));
                break;
            case "DSA":
                openSshKey = DSA_PREFIX + Base64.toBase64String(encodeDSA((DSAPublicKey) key));
                break;
            default:
                throw new IllegalArgumentException("Only RSA and DSA keys can be encoded to OpenSSH format by this helper");
        }

        return StringUtils.trim(openSshKey + " " + label);
    }

    private static byte[] encodeRSA(RSAPublicKey key) {
        try (ByteArrayOutputStream buf = new ByteArrayOutputStream()) {
            byte[] name = RSA_TYPE.getBytes("ASCII");
            writeByteArray(name, buf);
            writeByteArray(key.getPublicExponent().toByteArray(), buf);
            writeByteArray(key.getModulus().toByteArray(), buf);
            return buf.toByteArray();
        } catch (IOException e) {
            throw new KeyCodecException(e);
        }
    }

    private static byte[] encodeDSA(DSAPublicKey key) {
        try (ByteArrayOutputStream buf = new ByteArrayOutputStream()) {
            byte[] name = DSA_TYPE.getBytes("ASCII");
            writeByteArray(name, buf);
            writeByteArray(key.getParams().getP().toByteArray(), buf);
            writeByteArray(key.getParams().getQ().toByteArray(), buf);
            writeByteArray(key.getParams().getG().toByteArray(), buf);
            writeByteArray(key.getY().toByteArray(), buf);
            return buf.toByteArray();
        } catch (IOException e) {
            throw new KeyCodecException(e);
        }
    }

    private static void writeByteArray(byte[] str, OutputStream os) throws IOException {
        for (int shift = 24; shift >= 0; shift -= 8) {
            os.write((str.length >>> shift) & 0xFF);
        }
        os.write(str);
    }
}