import java.io.*;
import java.security.*;
import java.util.Calendar;


/**
 * @author Kalam Shah
 * @version 1.0 September 1 2003
 * This class calculates Digests for SHA-1 and MD5
 * It also signs using DSA, and RSA.
 * Finally, it also verifies the signs
 */
public class MdRsaDsa {
    
    //The DSA/RSA public and private keys along with the signature
    private PublicKey dsaKu;
    private PrivateKey dsaKr;
    private Signature dsaSign;
    
    public java.security.PrivateKey getDsaKr() {
        return dsaKr;
    }
    public void setDsaKr(java.security.PrivateKey dsaKr) {
        this.dsaKr = dsaKr;
    }
    
    public java.security.PublicKey getDsaKu() {
        return dsaKu;
    }
    public void setDsaKu(java.security.PublicKey dsaKu) {
        this.dsaKu = dsaKu;
    }
    
    public java.security.Signature getDsaSign() {
        return dsaSign;
    }
    public void setDsaSign(java.security.Signature dsaSign) {
        this.dsaSign = dsaSign;
    }

    /**
     * Though this function is named dsaSign, it is used for both RSA as well]
     * as DSA signing
     *
     *@param filename: the file which is used as input for signing
     *@param keySize: The key size to be used in bits
     *@param msgSize: The message size in bytes
     *@param algo: Can be either RSA or DSA
     *@param hash: Can be either 'RSAwithSHA1', 'RSAwithMD5' or 'DSAwithSHA1'
     */
    public byte[] dsaSign(String inFile, String outFile, int keySize, int msgSize,
                                                         String algo, String hash) {
        try {
            /************************* BEGIN Generate a key pair ***************************/
            long startTime = Calendar.getInstance().getTimeInMillis();
            KeyPairGenerator keygen = KeyPairGenerator.getInstance(algo);
            SecureRandom random = SecureRandom.getInstance("SHA1PRNG","SUN");
            
            keygen.initialize(keySize,random);

            KeyPair keypair = keygen.generateKeyPair();
            this.dsaKr = keypair.getPrivate();
            this.dsaKu = keypair.getPublic();
            /************************* END Generate a key pair ***************************/
            
            long midTime = Calendar.getInstance().getTimeInMillis();

            /************* Create a Signature object and initialize it with private key **************/
            this.dsaSign = Signature.getInstance(hash);
            this.dsaSign.initSign(this.dsaKr);

            /************************* Update and Sign the data ***************************/
            FileInputStream fis = new FileInputStream(inFile);
            BufferedInputStream bis = new BufferedInputStream(fis);
            byte buffer[] = new byte[64]; //512 bits=64 bytes and SHA-1/MD5 uses 512 bit blocks
            int length;
            int totalLength = 0;
            while(bis.available() != 0 && totalLength<msgSize){
                    length = bis.read(buffer);
                    totalLength += length;
                    this.dsaSign.update(buffer,0,length);
            }
            bis.close();

            /*
             * Now, that all the data to be signed is read in,
             * and hash generated ( using update method )
             * generate a signature for it
             */
            byte[] realsign = this.dsaSign.sign();
            long endTime = Calendar.getInstance().getTimeInMillis();

            /*********** Save statistics in file**********/
            FileWriter fw = new FileWriter(outFile, true);
            fw.write("\n\n\n************* "+hash+" ***************\n");
            fw.write("Message Size:"+totalLength+" bytes       KeySize:"+keySize+" bits\n");
            fw.write("********Key generation time = "+(midTime - startTime)+" ms\n");
            fw.write("********Signing time = "+(endTime - midTime)+" ms\n");
            fw.close();
            return (realsign);
        }catch(Exception e) {
            System.out.println("ERROR:"+e);
            return null;
        }
    }
    
    /**
     * Though this function is named dsaVerify, it is used for both DSA and RSA
     * verification.
     *
     *@param inFile: Input file name
     *@param outFile: Output file name
     *@param realsign: The sign which is to be veirfied
     *@param msgSize: The message size under consideration
     *@param keySize: The keysize used for signing
     */
    public boolean dsaVerify (String inFile, String outFile, byte[] realsign, int msgSize, int keySize) {
        try {
            long startTime = Calendar.getInstance().getTimeInMillis();
            this.dsaSign.initVerify(this.getDsaKu());
            
            //read the input
            FileInputStream fis = new FileInputStream(inFile);
            BufferedInputStream bis = new BufferedInputStream(fis);
            byte buffer[] = new byte[64]; //512 bits=64 bytes and SHA-1/MD5 uses 512 bit blocks
            int length;
            int totalLength = 0;
            while(bis.available() != 0 && totalLength<msgSize){
                    length = bis.read(buffer);
                    totalLength += length;
                    this.dsaSign.update(buffer,0,length);
            }
            bis.close();
            
            //verify
            boolean verified = this.dsaSign.verify(realsign);
            
            long endTime = Calendar.getInstance().getTimeInMillis();
            
            //write result in output file
            FileWriter fw = new FileWriter(outFile, true);
            fw.write("********Time taken to verify: "+(endTime-startTime)+" ms\n");
            
            fw.close();
            
            return (verified);
        }catch(Exception e) {
            System.out.println("ERROR:"+e);
            return false;
        }
    }
    
    /**
     * This function is used for making the message digest using MD5 or SHA-1
     *
     *@param inFile: The input file for used as a message
     *@param outFile: The file in which the results will be appended
     *@param algo: Can be either MD5 or SHA-1
     *@param msgSize: The size of message being considered for calculating the
     *                digest
     */
    public void makeDigest(String inFile, String outFile, String algo, int msgSize) {
        try {
            long startTime = Calendar.getInstance().getTimeInMillis();
            
            MessageDigest objMsgDig = MessageDigest.getInstance(algo);
            
            FileInputStream fis = new FileInputStream(inFile);
            BufferedInputStream bis = new BufferedInputStream(fis);
            
            /**
             * Each block should be in multiples of 512 bits = 64 bytes
             */
            byte buffer[] = new byte[64];
            //read the input
            int length;
            int totalLength = 0;
            while(bis.available() != 0 && totalLength<msgSize){
                    length = bis.read(buffer);
                    totalLength += length;
                    objMsgDig.update(buffer,0,length);
            }
            bis.close();
            
            //calculate the digest
            byte [] realdigest = objMsgDig.digest();
            
            long endTime = Calendar.getInstance().getTimeInMillis();

            //write output in a file
            FileWriter fw = new FileWriter(outFile, true);
            fw.write("\n\n\n************ "+algo+" *************\n");
            fw.write("Message size: "+totalLength+" bytes\n");
            fw.write("********Time taken to verify: "+(endTime-startTime)+" ms\n");
            fw.close();
        }catch(Exception e) {
            System.out.println("ERROR:"+e);
        }
    }
    
    public static void main(String [] args) {
        if(args.length!=1) {
            System.out.println("Usage: DSASignGen <name_of_file_to_sign>");
            System.exit(0);
        }
        
        //make object
        MdRsaDsa objMdRsaDsa = new MdRsaDsa();
        
        //make keysize, messageSize and algo arrays
        int arrKeySize[] = {512,768,1024};
        int arrMsgSize[] = {512, 1024, 2048, 4096, 8192};
        String arrSignAlgo[] = {"DSA", "SHA1withDSA", "RSA", "SHA1withRSA", "RSA", "MD5withRSA"};
        
        String inFile = args[0];
        String outFile = args[0]+".results";
        
        //call the DSA/RSA algorithm in a loop
        byte[] temp = null;
        for (int ai=0; ai<6; ai+=2) {
            for (int i=0; i<arrKeySize.length; i++) {
                for (int j=0; j<arrMsgSize.length; j++) {
                    //sign with DSA
                    System.out.println(arrSignAlgo[ai+1]+" Algorithm: keySize="+arrKeySize[i]+" bits    "
                            + "Message Size="+arrMsgSize[j]+" bytes");
                    temp = objMdRsaDsa.dsaSign(inFile, outFile, arrKeySize[i], arrMsgSize[j], 
                                                    arrSignAlgo[ai],arrSignAlgo[ai+1]);
                    System.out.print("Signed.");
                    //verify the sign
                    if (objMdRsaDsa.dsaVerify(inFile, outFile, temp, arrMsgSize[j], arrKeySize[i]))
                        System.out.println("Verified");
                    else System.out.println("Verification Failed.");
                }
            }
        }
        
        //For the message digests: MD5 and SHA-1
        String arrDigest[] = {"SHA-1","MD5"};
        for (int ai=0; ai<arrDigest.length; ai++) {
            for (int i=0; i<arrMsgSize.length; i++) {
                System.out.println(arrDigest[ai]+" Algorithm: Message Size: "+arrMsgSize[i]+" bytes");
                objMdRsaDsa.makeDigest(inFile, outFile, arrDigest[ai], arrMsgSize[i]);
            }
        }
        
        System.out.println("Completed.\nPlease check output file:"+outFile);
    }
    
    
    
}