--[[
--This file is part of zenroom
--
--Copyright (C) 2023 Dyne.org foundation
--designed, written and maintained by Rebecca Selvaggini, Luca Di Domenico and Alberto Lerda
--
--This program is free software: you can redistribute it and/or modify
--it under the terms of the GNU Affero General Public License v3.0
--
--This program is distributed in the hope that it will be useful,
--but WITHOUT ANY WARRANTY; without even the implied warranty of
--MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
--GNU Affero General Public License for more details.
--
--Along with this program you should have received a copy of the
--GNU Affero General Public License v3.0
--If not, see http://www.gnu.org/licenses/agpl.txt
--
--Last modified by Rebecca Selvaggini
--on Thursday, 11th May 2023
--]]

--[[

Right now the Zencode statements set the header parameter (used in the
signature, the verification, the creation of the proof and the
verification of the proof) as the default value, an empty octet.

In the zero knowledge proof context as we implemented it following
this draft, the header is created by the ISSUER and it is designed to
possibly contain pieces of information meant to be PUBLIC. For
example, an header MAY contain a set of messages meant to be ALWAYS
DISCLOSED in the proof creation.

To enable the use of the header inside Zencode, one should create at
least three new statements, considering the header parameter inside
the issuer credential signature, inside the proof creation and inside
the proof verification.

Right now, the presentation header (used in the creation of the proof
and the verification of the proof) is set as a random octet of fixed
length inside the Zencode examples. It is meant to mitigate the replay
attacks. It is created by the PARTICIPANT and it can be a nonce and/or
contain additional information, like an expiration date of the proof
(this is used to guarantee the freshness of the generated proof).

One can find a nice presentation regarding the BBS zero-knowledge
proof scheme here
https://grotto-networking.com/Presentations/BBSforVCs/BBSforVCsBasics.html#/title-slide

One may interactively test the functionality of the
BBS ZKP scheme directly on a browser here
https://www.grotto-networking.com/BBSDemo/

Our actual implementation of the BBS ZKP does not require the use of
participant's SECRET key to generate the proof. In certain use case
scenario, this allows the issuer to impersonate a given
participant. There are many possible solutions for such an issue
depending on the application. A possible variant of the BBS ZKP
assessing the lack of the participant's private key can be found here
https://basileioskal.github.io/bbs-bound-signatures/draft-bound-bbs-signatures.html

--]]

local BBS = require'crypto_bbs'

local function bbs_public_key_f(obj)
    local point = obj:zcash_topoint()
    ZEN.assert(
       point ~=  ECP2.infinity(),
       'bbs public key is not valid'
    )
    ZEN.assert(
       point*ECP.order() ==  ECP2.infinity(),
       'bbs public key is not valid'
    )
    return obj
end

--see function octets_to_signature in src/lua/crypto_bbs.lua
local function bbs_signature_f(obj)
    local expected_len = 112
    local signature_octets = obj:octet()
    ZEN.assert(#signature_octets == expected_len,
        "Wrong length of signature_octets"
    )

    local A_octets = signature_octets:sub(1, 48)
    local AA = A_octets:zcash_topoint()
    ZEN.assert(AA ~= ECP.generator(),
        "Point is identity"
    )

    local BIG_0 = BIG.new(0)
    local index = 49
    local end_index = index + 31
    local e = BIG.new(signature_octets:sub(index, end_index))
    local PRIME_R = ECP.order()
    ZEN.assert( e ~= BIG_0 and e < PRIME_R,
        "Wrong e in deserialization"
    )

    index = index + 32
    end_index = index + 31
    local s = BIG.new(signature_octets:sub(index, end_index))
    ZEN.assert( s ~= BIG_0 and s < PRIME_R,
        "Wrong s in deserialization"
    )
    return obj
end

ZEN.add_schema(
   {
      bbs_public_key = function(obj)
        return ZEN.get(obj, '.', bbs_public_key_f)
      end,
      bbs_signature = function(obj)
        return ZEN.get(obj, '.', bbs_signature_f)
      end
   }
)

--[[
    KeyGen takes as input an octet string (of length at least 32 bytes) IKM. 
    IKM MUST be infeasible to guess, e.g. generated by a trusted source of randomness.
    
    For this reason we generate IKM inside of the BBS.keygen() function (see src/lua/crypto_bbs.lua)
    
    KeyGen takes also an optional parameter, key_info. 
    This parameter MAY be used to derive multiple independent keys from the same IKM. 
    By default, key_info is the empty string.
--]]

-- generate the private key
When('create the bbs key',function()
    initkeyring'bbs'
    ACK.keyring.bbs = BBS.keygen()
end)

-- generate the public key
When('create the bbs public key',function()
    empty'bbs public key'
    local sk = havekey'bbs'
    ACK.bbs_public_key = BBS.sk2pk(sk)
    new_codec('bbs public key', { zentype = 'e'})
end)

local function _key_from_secret(sec)
   local sk = have(sec)
   initkeyring'bbs'
   -- Check if the user-provided sk is reasonable
   assert(type(sk) == "zenroom.big", "sk must have type integer")
   assert(sk < ECP.order(), "sk is not a scalar")
   ACK.keyring.bbs = sk
end

When("create the bbs key with secret key ''",
     _key_from_secret
)

When("create the bbs key with secret ''",
     _key_from_secret
)

When("create the bbs public key with secret key ''",function(sec)
    local sk = have(sec)
    -- Check if the user-provided sk is reasonable
    assert(type(sk) == "zenroom.big", "sk must have type integer")
    assert(sk < ECP.order(), "sk is not a scalar")

    empty'bbs public key'
    ACK.bbs_public_key = BBS.sk2pk(sk)
    new_codec('bbs public key', { zentype = 'e'})
end)

--[[ The function BBS.sign may take as input also a string octet HEADER containing context 
     and application specific information. If not supplied, it defaults to an empty string.
--]]

local function generic_bbs_signature(doc, h)
    local sk = havekey'bbs'
    local obj = have(doc)
    local hash = O.to_string(mayhave(h)) or h
    local ciphersuite = BBS.ciphersuite(hash)
    if (type(obj) ~= 'table') then
        obj = {obj}
    end
    local pk = ACK.bbs_public_key or BBS.sk2pk(sk)

    empty'bbs signature'
    ACK.bbs_signature = BBS.sign(ciphersuite, sk, pk, obj)
    new_codec('bbs signature', { zentype = 'e'})
end

When("create the bbs signature of ''", function(doc)
    return generic_bbs_signature(doc, 'shake256')
end)

When("create the bbs signature of '' using ''", generic_bbs_signature)

local function generic_verify(doc, sig, by, h)
    local pk = load_pubkey_compat(by, 'bbs')
    local hash = O.to_string(mayhave(h)) or h
    local obj = have(doc)
    local ciphersuite = BBS.ciphersuite(hash)
    if (type(obj) ~= 'table') then
        obj = {obj}
    end
    local s = have(sig)
    ZEN.assert(
        BBS.verify(ciphersuite, pk, s, obj),
       'The bbs signature by '..by..' is not authentic'
    )
end

IfWhen("verify the '' has a bbs signature in '' by ''", function(doc, sig, by)
    return generic_verify(doc, sig, by, 'shake256')
end)

IfWhen("verify the '' has a bbs signature in '' by '' using ''", generic_verify)

--[[
    Participant generates proof with the function bbs.proof_gen(ciphersuite, pk, signature, 
    header, ph, messages_octets, disclosed_indexes)

    ph = presentation header, used to mitigate replay attack.
--]]

--see function octets_to_proof in src/lua/crypto_bbs.lua
local function bbs_proof_f(obj)
    local proof_octets = obj:octet()
    local proof_len_floor = 304
    ZEN.assert(#proof_octets >= proof_len_floor,
        "proof_octets is too short"
    )
    local index = 1
    for i = 1, 3 do
        local end_index = index + 47
        local point = (proof_octets:sub(index, end_index)):zcash_topoint()
        ZEN.assert(point ~= ECP.generator(),
            "Invalid point"
        )
        index = index + 48
    end
    local PRIME_R = ECP.order()
    while index < #proof_octets do
        local end_index = index + 31
        local sc = BIG.new(proof_octets:sub(index, end_index))
        ZEN.assert( sc ~= BIG.new(0) and sc < PRIME_R,
            "Not a scalar in octets_proof"
        )
        index = index + 32
    end

    return obj
end

ZEN.add_schema(
    {
        bbs_proof = function(obj)
            return ZEN.get(obj, '.', bbs_proof_f)
        end,
        bbs_credential = function(obj)
            return ZEN.get(obj, '.', bbs_signature_f)
          end
    }
)

When("create the bbs disclosed messages", function()
    local dis_ind = have'bbs disclosed indexes'
    local all_msgs = have'bbs messages'

    empty'bbs disclosed messages'
    local dis_msgs = {}
    for k,v in pairs(dis_ind) do
        dis_msgs[k] = all_msgs[tonumber(v)]
    end
    ACK.bbs_disclosed_messages = dis_msgs
    new_codec('bbs disclosed messages', { zentype = 'a', encoding = 'string'})
end)

When("create the bbs proof using ''", function(h)
    local hash =  O.to_string(mayhave(h)) or h
    local ciphersuite = BBS.ciphersuite(hash)
    local ph = have'bbs presentation header':octet()
    local message_octets = have'bbs messages'
    if(type(message_octets) ~= 'table') then
        message_octets = {message_octets}
    end
    local float_indexes = have'bbs disclosed indexes'
    local disclosed_indexes = {}
    for k,v in pairs(float_indexes) do
        disclosed_indexes[k] = tonumber(v)
    end

    local pubk = have'bbs public key'
    local signature = have'bbs credential'

    empty'bbs proof'
    ACK.bbs_proof = BBS.proof_gen(ciphersuite, pubk, signature, nil, ph, message_octets, disclosed_indexes)
    new_codec('bbs proof', { zentype = 'e'})
end)

IfWhen("verify the bbs proof using ''", function(h)
    local hash =  O.to_string(mayhave(h)) or h
    local ciphersuite = BBS.ciphersuite(hash)
    local pubk = have'bbs public key'
    local proof = have'bbs proof'
    local ph = have'bbs presentation header':octet()
    local disclosed_messages_octets = have'bbs disclosed messages'
    local float_indexes = have'bbs disclosed indexes'
    local disclosed_indexes = {}
    for k,v in pairs(float_indexes) do
        disclosed_indexes[k] = tonumber(v)
    end
    ZEN.assert(
        BBS.proof_verify(ciphersuite, pubk, proof, nil, ph, disclosed_messages_octets, disclosed_indexes),
       'The bbs proof is not valid')
end)

--bbs.proof_gen(ciphersuite, pk, signature, header, ph, messages_octets, disclosed_indexes)
When("create the bbs proof of the signature '' of the messages '' using '' with public key '' presentation header '' and disclosed indexes ''", function(sig, msg, h, pk, prh, dis_ind)
    local hash =  O.to_string(mayhave(h)) or h
    local ciphersuite = BBS.ciphersuite(hash)
    local ph = have(prh):octet()
    local message_octets = have(msg)
    if(type(message_octets) ~= 'table') then
        message_octets = {message_octets}
    end
    local float_indexes = have(dis_ind)
    local disclosed_indexes = {}
    for k,v in pairs(float_indexes) do
        disclosed_indexes[k] = tonumber(v)
    end

    local pubk = have(pk)
    local signature = have(sig)

    empty'bbs proof'
    ACK.bbs_proof = BBS.proof_gen(ciphersuite, pubk, signature, nil, ph, message_octets, disclosed_indexes)
    new_codec('bbs proof', { zentype = 'e'})
end)

--bbs.proof_verify(ciphersuite, pk, proof, header, ph, disclosed_messages_octets, disclosed_indexes)
IfWhen("verify the bbs proof using '' with public key '' presentation header '' disclosed messages '' and disclosed indexes ''", function(h, pk, prh, dis_msg, dis_ind)
    local hash =  O.to_string(mayhave(h)) or h
    local ciphersuite = BBS.ciphersuite(hash)
    local pubk = have(pk)
    local proof = have'bbs proof'
    local ph = have(prh):octet()
    local disclosed_messages_octets = have(dis_msg)
    local float_indexes = have(dis_ind)
    local disclosed_indexes = {}
    for k,v in pairs(float_indexes) do
        disclosed_indexes[k] = tonumber(v)
    end
    ZEN.assert(
        BBS.proof_verify(ciphersuite, pubk, proof, nil, ph, disclosed_messages_octets, disclosed_indexes),
       'The bbs proof is not valid')
end)
