/* Group Master
////////////////////
    Shamir_IBS ibs = new Shamir_IBS(2048);
    BigInteger n = ibs.getN();
    BigInteger e = ibs.getE();
    BigInteger g = ibs.deriveKey("user");
*/


/* User
////////////////////
    Shamir_IBS ibs = new Shamir_IBS(n,e,g);
    Pair<byte[], byte[]> signature = ibs.sign("Hello world!");
    boolean verificationResult = ibs.verify("user", "Hello world!", signature);
*/


import android.util.Pair;
import java.math.BigInteger;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.Random;

public class Shamir_IBS {

    private BigInteger p;
    private BigInteger q;
    private BigInteger n;
    private BigInteger phi_n;
    private BigInteger e;
    private BigInteger d;
    private BigInteger g;

    public Shamir_IBS(){}

    public Shamir_IBS(int bitLength)
    {
        Random rnd = new Random();
        p = BigInteger.probablePrime(bitLength/2,rnd);
        q = BigInteger.probablePrime(bitLength/2,rnd);

        n = p.multiply(q);
        phi_n = p.subtract(BigInteger.ONE).multiply(q.subtract(BigInteger.ONE));

        do {
            rnd = new Random();
            e = new BigInteger(256,rnd).mod(phi_n);
        } while(e.gcd(phi_n).equals(BigInteger.ONE) == false);

        d = e.modInverse(phi_n);
    }

    public BigInteger getN() {
        return n;
    }

    public BigInteger getE() {
        return e;
    }

    public Shamir_IBS(BigInteger n, BigInteger e, BigInteger g)
    {
        this.n = n;
        this.e = e;
        this.g = g;
    }

    public BigInteger deriveKey(String ID)  throws NoSuchAlgorithmException
    {
        MessageDigest myDigest = MessageDigest.getInstance("SHA-256");
        g = new BigInteger(myDigest.digest(ID.getBytes())).modPow(d,n);
        return g;
    }

    public Pair<byte[], byte[]> sign(String mes) throws NoSuchAlgorithmException {
        Random rnd = new Random();
        BigInteger r = new BigInteger(n.bitLength(),rnd);
        BigInteger t = r.modPow(e,n);

        MessageDigest digest = MessageDigest.getInstance("SHA-256");
        byte[] message_hash = digest.digest(mes.concat(t.toString()).getBytes());
        BigInteger hm = new BigInteger(message_hash);

        BigInteger s = r.modPow(hm,n).multiply(g).mod(n);

        /* signature is (s,t) */
        /* Remove the sign bit */
        byte[] ArrayFirst = s.toByteArray();
        if (ArrayFirst[0] == 0) {
            byte[] tmp = new byte[ArrayFirst.length - 1];
            System.arraycopy(ArrayFirst, 1, tmp, 0, tmp.length);
            ArrayFirst = tmp;
        }
        byte[] ArraySecond = t.toByteArray();
        if (ArraySecond[0] == 0) {
            byte[] tmp = new byte[ArraySecond.length - 1];
            System.arraycopy(ArraySecond, 1, tmp, 0, tmp.length);
            ArraySecond = tmp;
        }
        return new Pair<byte[], byte[]>(ArrayFirst, ArraySecond);
    }


    public boolean verify(String ID, String mes, Pair<byte[], byte[]> sig) throws NoSuchAlgorithmException {
        /* signature is (s,t) */
        /* verification */
        BigInteger s = new BigInteger(ConcatArrays(new byte[]{0},sig.first));
        BigInteger t = new BigInteger(ConcatArrays(new byte[]{0},sig.second));

        BigInteger ss = s.modPow(e,n);

        MessageDigest digest = MessageDigest.getInstance("SHA-256");
        byte[] message_hash = digest.digest(mes.concat(t.toString()).getBytes());
        BigInteger hm = new BigInteger(message_hash);

        digest = MessageDigest.getInstance("SHA-256");
        byte[] id_hash = digest.digest(ID.getBytes());
        BigInteger hID = new BigInteger(id_hash);

        BigInteger tt = t.modPow(hm,n).multiply(hID).mod(n);

        if (ss.equals(tt)) return true;
        else return false;
    }


    private static byte[] ConcatArrays(byte[] first, byte[] second) {
        int totalLength = first.length + second.length;
        byte[] result = Arrays.copyOf(first, totalLength);
        int offset = first.length;
        System.arraycopy(second, 0, result, offset, second.length);
        return result;
    }
}