/*****************************************************************************
 *
 * Copyright (c) 2018-19, Josef Treus
 * Westfälische Wilhelms-Universität Münster, AG Embedded Systems
 * All rights reserved.
 *
 * This file is part of SysCSec
 *
 * SysCSec is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your
 * option) any later version.
 *
 * SysCSec 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 General Public License
 * for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with SysCSec.  If not, see <http://www.gnu.org/licenses/>.
 *
 *
 *  Please report any problems or bugs to: joseftreus@web.de
 *
 ****************************************************************************/
#ifndef ZEMD_ATTACKER_CPP
#define ZEMD_ATTACKER_CPP

#include "rsa.cpp"
#include "parameters.h"
#include "Utils.cpp"
#include "RessourceHandler.cpp"


/* The attacker receives the public key from the subject and starts a ZEMD attack.
 * Param: sc_module_name name : name of the attacker
 *        int msgCount : number of messages used in each step
*/
SC_MODULE(ZEMD_Attacker) {
    sc_fifo_in<bigint>in;
    sc_fifo_out<bigint>out;
    sc_in<bool>cipherProceeded_in;

    bigint publicKeyE;
    bigint publicKeyN;

    int msgCount;
    int roundsPerBit;

    /* ZEMD attack
     * Param: -
     * Return: -
     */
    void zemd() {
        // wait for the public key
        wait(in->data_written_event());
        this->publicKeyE=in->read();
        this->publicKeyN=in->read();
        // determine lentgh of key (assume d = n)
        int n= publicKeyN.to_int();
        int bitsThatAreZero= Utils::unsignificantBitCount(n, NBITS);
        int keyLength=NBITS-bitsThatAreZero;

        /* estimate key and set first bit to one
         * leading zeros are ignored
         */
        bool keyBitGuesses[keyLength]= {0};

        // guess one bit after the other
        keyBitGuesses[keyLength-1]=true;
        for(int i = keyLength-2; i >= 0; i--) {
            cout << "Determine bit "<< i << endl;
            int zeroConcluded,oneConcluded;
            for(int h = 0; h < roundsPerBit; h++){
                // try with one and zero
                map<int,int> D_Zero, D_One;
                for(int j=0; j < 2; j++) {

                    // first guess: key bit i is one
                    if(j==1) {
                        keyBitGuesses[i]=true;
                    }
                        // second guess: key bit i is zero
                    else {
                        keyBitGuesses[i]=false;
                    }
                    int keyGuess = 0;
                    for(int l = keyLength-1; l >= i; l--){
                        keyGuess |= ((int)keyBitGuesses[l]) << l;
                    }
                    int meanHam=0; // mean Hamming weight

                    map<int,map<int, int>> S;
                    int decyptedMsg[msgCount]= {0};

                    char randVal[NBITS-bitsThatAreZero-1];
                    /* in each step, a random message is generated and
                        * encrypted with both the generated key and
                        * the original key.  The power trace of the
                        * original encryption is tracked and the
                        * Hamming weight of the result of the
                        * encryption with the generated key is
                        * determined.
                        */
                    for(int k=0; k < msgCount; k++) {
                        rand_bitstr(randVal, keyLength-1);
                        bigint randValAsBigint = randVal;
                        // decrption up to bit i - copied and modified from rsa.cpp
                        bigint d = 1;
                        for ( int l = keyLength-1; l >=0; l-- )
                        {
                            d = ( d * d ) % publicKeyN;
                            if ( keyBitGuesses[l] ){
                            
                            d = ( d * randValAsBigint ) % publicKeyN;
                            }
                        }
                        decyptedMsg[k] = d.to_int();  
                        meanHam += Utils::hammingWeight(decyptedMsg[k]);
                        out->write(randValAsBigint);
                        wait();
                        S[k]= RessourceHandler::getInstance()->powerTraces[RessourceHandler::getInstance()->powerTraces.size()-1];
                    }
                    meanHam /= msgCount;
                    /* The power traces are split into two subsets according to the Hamming weight, with the calculated mean as threshold.
                        * Traces with exactly mean Hamming weight are skipped.
                        */
                    map<int, map<int,int>> sHigh;
                    map<int, map<int, int>> sLow;
                    
                    for(int k=0; k < msgCount; k++) {
                        if(Utils::hammingWeight(decyptedMsg[k]) > meanHam) {
                            sHigh[sHigh.size()]=S[k];
                        }else if(Utils::hammingWeight(decyptedMsg[k]) < meanHam) {
                            sLow[sLow.size()]=S[k];
                        }
                    }
                    
                    // compute average trace
                    map<int, int>_sHigh = RessourceHandler::getInstance()->meanTrace(sHigh);
                    map<int, int>_sLow = RessourceHandler::getInstance()->meanTrace(sLow);
                    map<int,int> D = RessourceHandler::getInstance()->traceDifference(_sHigh, _sLow);
                    // compute difference trace
                    if(j == 0){
                        D_Zero = D;
                    }else{
                        D_One = D;
                    }
                    RessourceHandler::getInstance()->additionalTraces[((keyLength-2)*2*roundsPerBit) - (i*2*roundsPerBit)+j+2*h] = D;
                }
                /* Determine and sum up the variances of all difference traces of both zero and one bit guess.
                */
                double abweichungZero = 0;
                double abweichungOne = 0;
                abweichungZero = RessourceHandler::getInstance()->traceVariance(D_Zero);
                abweichungOne = RessourceHandler::getInstance()->traceVariance(D_One);
                /* Depending on which difference trace has a greater variance and thus greater peaks
                * the final estimate is a zero or a one.
                */
                if(abweichungZero < abweichungOne){
                    zeroConcluded++;
                }else{
                    oneConcluded++;
                }
            }
            if(zeroConcluded < oneConcluded){
                keyBitGuesses[i] = false;
            }else{
                keyBitGuesses[i] = true;
            }
        }
       
        cout << endl;
        for(int r= keyLength-1; r >= 0; r--) {
            if(keyBitGuesses[r]) {
                cout << 1;
            }

            else {
                cout << 0;
            }
        }
        cout << endl;
        RessourceHandler::getInstance()->powerTraces.clear();
        wait(1, SC_SEC);
        sc_stop();
    }

    SC_HAS_PROCESS(ZEMD_Attacker);
	ZEMD_Attacker(sc_module_name name, int msgCount, int roundsPerBit): sc_module(name), cipherProceeded_in("cPi"), in("in"), out("out"){
        this->msgCount = msgCount;
        this->roundsPerBit = roundsPerBit;
        SC_THREAD(zemd);
        sensitive << cipherProceeded_in.pos() << cipherProceeded_in.neg();
    }
};
#endif
