package com.jscape.inet.ssh.protocol.v2.marshaling.keys;

import com.jscape.inet.ssh.protocol.v2.messages.PublicKeyAlgorithm;
import com.jscape.inet.ssh.protocol.v2.messages.Signature;
import com.jscape.util.ArrayTools;

import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAKey;

/**
 * @author Alex Usun
 */
public class JceRsaSigner
        implements Signer {

    private final Object provider;

    public JceRsaSigner(Object provider) {
        this.provider = provider;
    }

    @Override
    public Signature sign(byte[] data, String publicKeyAlgorithm, PrivateKey key)
            throws Exception {
        java.security.Signature signature = signatureFor(publicKeyAlgorithm);
        signature.initSign(key);
        signature.update(data);
        byte[] blob = signature.sign();

        return new Signature(publicKeyAlgorithm, blob);
    }

    @Override
    @SuppressWarnings("unchecked")
    public void assertSignatureValid(byte[] data, PublicKey key, Signature signatureData)
            throws Exception {
        java.security.Signature signature = signatureFor(signatureData.identifier);
        signature.initVerify(key);
        signature.update(data);

        byte[] signatureBlob = adjust(signatureData.blob, key);

        if (!signature.verify(signatureBlob)) {
            throw new InvalidSignatureException();
        }
    }

    private java.security.Signature signatureFor(String publicKeyAlgorithm)
            throws Exception {
        String signatureAlgorithm = String.format("%sWithRSA", PublicKeyAlgorithm.signatureAlgorithmFor(publicKeyAlgorithm));
        return Signatures.signatureFor(signatureAlgorithm, this.provider);
    }

    public byte[] adjust(byte[] signatureBlob, PublicKey key) {
        if (key instanceof RSAKey) {
            RSAKey rsaKey = (RSAKey) key;
            int modulusLength = rsaKey.getModulus().bitLength();
            int blobLength = signatureBlob.length * 8;

            if (blobLength < modulusLength) {
                int paddingBytes = (modulusLength - blobLength) / 8;
                return ArrayTools.copy(signatureBlob, 0, new byte[signatureBlob.length + paddingBytes], paddingBytes, signatureBlob.length);
            }
        }

        return signatureBlob;
    }

}
