%maven org.cryptimeleon:math:1.0.0
%maven org.cryptimeleon:craco:1.0.0

import org.cryptimeleon.math.structures.groups.elliptic.nopairing.Secp256k1;
import org.cryptimeleon.math.structures.groups.lazy.*;
import org.cryptimeleon.math.hash.impl.SHA256HashFunction;
import org.cryptimeleon.math.structures.rings.zn.Zn;
import org.cryptimeleon.math.hash.impl.ByteArrayAccumulator;
import org.cryptimeleon.math.structures.groups.*;
import org.cryptimeleon.math.serialization.converter.JSONPrettyConverter;
import org.cryptimeleon.craco.protocols.arguments.sigma.schnorr.*;
import org.cryptimeleon.craco.protocols.*;
import org.cryptimeleon.craco.protocols.arguments.fiatshamir.FiatShamirProofSystem;

//Set up group and generate key
var group = new LazyGroup(new Secp256k1()); 
var H1 = new HashIntoLazyGroup(new Secp256k1.HashIntoSecp256k1(), group);
var H2 = new SHA256HashFunction();
var jsonConverter = new JSONPrettyConverter(); //for serialization later

var g = group.getGenerator();
var k = group.getUniformlyRandomNonzeroExponent(); //secret key
var h = g.pow(k).precomputePow(); //public key

byte[] evaluatePRF(Zn.ZnElement k, byte[] x) {
    var h2Preimage = new ByteArrayAccumulator();
    h2Preimage.escapeAndSeparate(h);
    h2Preimage.escapeAndSeparate(x);
    h2Preimage.escapeAndAppend(H1.hash(x).pow(k));
    
    return H2.hash(h2Preimage.extractBytes());
}

var result = evaluatePRF(k, new byte[] {4, 8, 15, 16, 23, 42});
Arrays.toString(result)

class ProofCommonInput implements CommonInput {
    public final GroupElement a,b;
    
    public ProofCommonInput(GroupElement a, GroupElement b) {
        this.a = a;
        this.b = b;
    }
}

class ProofWitnessInput implements SecretInput {
    public final Zn.ZnElement k;
    
    public ProofWitnessInput(Zn.ZnElement k) {
        this.k = k;
    }
}

class ReplyCorrectnessProof extends DelegateProtocol {
    @Override
    protected SendThenDelegateFragment.SubprotocolSpec provideSubprotocolSpec(CommonInput commonInput, SendThenDelegateFragment.SubprotocolSpecBuilder builder) {
        //In this method, we define (for prover and verifier alike) what the proof parameters shall be.
        //We want to prove knowledge of a single variable, namely k. So we register "k" of type Zp with the builder.
        var kVar = builder.addZnVariable("k", group.getZn());
        //the result is a variable kVar that we will reference in the following.
        
        //Now we need to define what shall be proven. For this, we reference the knowledge variable created above.

        //statementToBeProven is an Expression "a^k = b" with variable k (not something that's computed here and now)
        var statementToBeProven = ((ProofCommonInput) commonInput).a.pow(kVar).isEqualTo(((ProofCommonInput) commonInput).b); 
        
        //With this expression format we tell the framework to add the fragment for "a^k = b" to the proof
        builder.addSubprotocol("replyCorrect", new LinearStatementFragment(statementToBeProven));

        //Similarly, we handle the other fragment
        builder.addSubprotocol("publicKeyCorrect", new LinearStatementFragment(g.pow(kVar).isEqualTo(h)));

        //Aaaand that's it :) - We have defined to prove knowledge of k such that a^k = b and g^k = h.
        return builder.build();
    }
    
    @Override
    protected SendThenDelegateFragment.ProverSpec provideProverSpecWithNoSendFirst(CommonInput commonInput, SecretInput secretInput, SendThenDelegateFragment.ProverSpecBuilder builder) {
        //Here, we need to set up which witnesses the issuer shall use, i.e. their secret key.
        //For every builder.addZnVariable() above, we must set the witness here (usually from secretInput)
        builder.putWitnessValue("k", ((ProofWitnessInput) secretInput).k);
        
        //That's it already, that's all the additional info the prover needs.
        return builder.build();
    }

    @Override
    public BigInteger getChallengeSpaceSize() {
        return group.size();
    }
}

//Wrap the sigma protocol into a Fiat-Shamir proof system
var fiatShamirProofSystem = new FiatShamirProofSystem(new ReplyCorrectnessProof());

//User's perspective

byte[] x = new byte[] {4, 8, 15, 16, 23, 42}; //want PRF(x)

var r = group.getUniformlyRandomExponent();
GroupElement a = H1.hash(x).pow(r);

//Send (serialized) group element
String messageOverTheWire = jsonConverter.serialize(a.getRepresentation());
messageOverTheWire

//Server's perspective

//Deserialize the request (also ensures the request is a valid group element)
GroupElement aServer = group.restoreElement(jsonConverter.deserialize(messageOverTheWire));

//Compute the response
var bServer = aServer.pow(k);

//Compute the proof
var proofServer = fiatShamirProofSystem.createProof(new ProofCommonInput(aServer,bServer), new ProofWitnessInput(k));

//Send response
var responseRepresentation = new org.cryptimeleon.math.serialization.ListRepresentation(bServer.getRepresentation(), proofServer.getRepresentation());
var responseOverTheWire = jsonConverter.serialize(responseRepresentation);
responseOverTheWire

//Deserialize stuff
var deserializedResponse = jsonConverter.deserialize(responseOverTheWire);
var b = group.restoreElement(deserializedResponse.list().get(0));

//Check proof
var commonInput = new ProofCommonInput(a,b); //set common input for proof
var proof = fiatShamirProofSystem.restoreProof(commonInput, deserializedResponse.list().get(1)); //deserialize proof
assert fiatShamirProofSystem.checkProof(commonInput, proof); //check proof

//Unblind b and compute PRF value
var h2Preimage = new ByteArrayAccumulator();
h2Preimage.escapeAndSeparate(h);
h2Preimage.escapeAndSeparate(x);
h2Preimage.escapeAndAppend(b.pow(r.inv()));

var result = H2.hash(h2Preimage.extractBytes());
Arrays.toString(result)