#include "lib/md5.h"
#include "lib/sha256.h"
#include <iostream>
#include <thread>
#include <string>
#include <fstream>
#include <vector>
#include <algorithm>
#include <chrono>
#include <cstring>
#include <atomic>
#include <unistd.h>
#include <random>
#include <sys/stat.h>
#include "sys/sysinfo.h"





using namespace std;
using namespace chrono;

#define PBSTR "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"
#define PBWIDTH 60

enum Error {
  ERR_UNKNOWN_OPTION, 
  ERR_CANT_OPEN_FILE,
  ERR_ARGS,
  ERR_THREADS
}; 

void error(Error n, string file)
{
  switch (n) {
   
    case ERR_UNKNOWN_OPTION:
      cout << "Error, unknown option" << endl;
      break;

    case ERR_CANT_OPEN_FILE:
      cout << "Error, can't open file: "<< file << endl;
      break;

    case ERR_ARGS:
      cout<< "Usage:\ncrack -f [file] -d [dict] -t [threads] --trace [worst performance] --check [sys info]" << endl;
      break;
    
    case ERR_THREADS:
      cout<< "Error, wrong number of threads" << endl;
      break;
  }
}

enum Info{
    CHECK_SYS
};

void info(Info i){
    struct sysinfo memInfo;
    sysinfo (&memInfo);
    long long totalPhysMemBytes = memInfo.totalram;
    long long totalPhysMemGBytes = (totalPhysMemBytes/(1024*1024*1024));
    const auto processor_count = std::thread::hardware_concurrency();

    switch (i) {

        case CHECK_SYS:
            cout << "Number of threads: " << processor_count << endl;
            cout << "Total RAM: " << totalPhysMemGBytes << "GB ("<<totalPhysMemBytes<<" bytes)"<<endl;
            break;
    }
}

void loadingFunction(double percentage){
    int val = (int) (percentage * 100);
    int lpad = (int) (percentage * PBWIDTH);
    int rpad = PBWIDTH - lpad;
    printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, "");
    fflush(stdout);
}

long long countLines(string &nombre){
    ifstream inFile(nombre); 
    long long c = count(istreambuf_iterator<char>(inFile), 
    istreambuf_iterator<char>(), '\n')+1;
    return c;
}

void readFile(string *&words,string &nombre, long long &size){
    ifstream fichero; string fila;
    words = new string[size];

    long long i = 0;
    fichero.open(nombre.c_str(), ios::in);
    if(fichero.is_open()){
        while(!fichero.eof()){
            getline(fichero, fila);
            words[i] = fila;
            i++;
        }
        fichero.close();
    }
}

void readFileDict(char *&dict,string &nombre, long long &sizeWordsDICT){
    cout<<"Loading dictionary in memory..."<<endl;
    std::ifstream fichero(nombre, ios::in);
    fichero.read(dict, sizeWordsDICT);
}

void crackPasswords(string *wordsMD5, char *&dict, string *&result, long long *&rangs, long long &sizeMD5, long long &sizeWordsDICT, int core, int &recovered, bool &trace)
{
    MD5 hashMD5;

    for(long long j=rangs[0];j<rangs[1];j++){
        if(core==0 && trace){
            loadingFunction(static_cast<double>(j)/static_cast<double>(rangs[1]));
        }
        string word = "";
        while(dict[j] != '\n' && j<=sizeWordsDICT){
            word = word + dict[j];
            j++;
        }
        word.erase(std::remove(word.begin(), word.end(), '\r' ), word.end());
        string wordHashMD5 = hashMD5(word);
        for(int word2crack=0;word2crack<sizeMD5;word2crack++){
            if(result[word2crack]!=""){
                continue;
            }
            if(wordsMD5[word2crack]==wordHashMD5){
                recovered++;
                result[word2crack] = word;
            }
        }
    }
    
}

void fillSplit(long long **&rangs, long long &sizeWordsDICT, int &numThreads){
    long long split = sizeWordsDICT/numThreads;
    long long resto = sizeWordsDICT%numThreads;
    long long endSPlit = split+resto;
    long long currentB=0;
    long long currentF=split-1;
    

    for(int rang=0; rang<numThreads; rang++){
        rangs[rang] = new long long[2];
        rangs[rang][0]=currentB;
        rangs[rang][1]=currentF;
        if(rang + 2 == numThreads){
            currentB = currentF+1;
            currentF += endSPlit;
        }
        else{
            currentB = currentF+1;
            currentF += split;
        }
    }
}

bool checkIfExist(string &nombre){
    ifstream fichero;
    fichero.open(nombre.c_str(), ios::in);
    if(fichero.is_open()){
        return true;
    }
    else{
        return false;
    }
}

string gen_random(const int len) { 
    static const char alphanum[] =
        "0123456789"
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "abcdefghijklmnopqrstuvwxyz";
    std::string tmp_s;
    tmp_s.reserve(len);

    random_device rd;
    mt19937 mt(rd());
    uniform_real_distribution<double> dist(0, sizeof(alphanum)-1);

    for (int i = 0; i < len; ++i) {
        tmp_s += alphanum[(int)(dist(mt))];
    }
    
    return tmp_s;
}

void outputFileCrack(string *result, long long &sizeMD5){
    ofstream fichero("plain.txt");
    if(fichero.is_open()){
        for(long long i =0;i<sizeMD5;i++){
            fichero<<result[i]<<endl;
        }
        fichero.close();
    }
}

void outputFileSHA256(string *result, long long &sizeMD5){
    srand((unsigned)time(NULL) * getpid());     
    string salt = gen_random(12);
    SHA256 sha256;
    ofstream fichero("new_passwords.txt");
    if(fichero.is_open()){
        for(long long i =0;i<sizeMD5;i++){
            if(result[i]==""){
               fichero<<""<<endl; 
            }
            else{
                string salt = gen_random(12);
                string resultString = result[i] + salt;
                fichero<<sha256(resultString)<<":"<<salt<<endl;
            }
        }
        fichero.close();
    }
}

long long GetFileSize(string filename)
{
    struct stat stat_buf;
    int rc = stat(filename.c_str(), &stat_buf);
    return rc == 0 ? stat_buf.st_size : -1;
}



int main(int argc,char *argv[])
{
    // create a new hashing object
    //SHA256 sha256;
    //std::cout << sha256("Hello World") << std::endl;

    //Parametros de entrada.
    string nameFileMD5="";
    string nameFileDict="";
    bool trace=false;
    int numThreads = 0;





    //Args
    int args = 0;
    while(args < argc){
        if(strcmp(argv[args], "--check") == 0){	
            info(CHECK_SYS);
            return -1;
        }
        if(strcmp(argv[args], "--help") == 0){	
            if(args+1<argc){
                error(ERR_ARGS,"");
                return -1;
            }
        }
        if(strcmp(argv[args], "-f") == 0){	
            if(args+1<argc){
                nameFileMD5=argv[args+1];
            }
        }
        if(strcmp(argv[args], "--trace") == 0){	
            trace=true;
        }
        if(strcmp(argv[args], "-d") == 0){	
            if(args+1<argc){
                nameFileDict=argv[args+1];
            }
        }
        if(strcmp(argv[args], "-t") == 0){	
            if(args+1<argc){
                numThreads=stoi(argv[args+1]);
                const auto processor_count = std::thread::hardware_concurrency();
                if(numThreads == 0 || processor_count<numThreads){
                    error(ERR_THREADS,"");
                    return -1;
                }
            }
        }
        args++;
    }
    if(nameFileMD5=="" || nameFileDict=="" || numThreads<=0){
        error(ERR_ARGS,"");
        return -1;
    }
    if(!checkIfExist(nameFileMD5)){
        error(ERR_CANT_OPEN_FILE, nameFileMD5);
        return -1;
    }
    if(!checkIfExist(nameFileDict)){
        error(ERR_CANT_OPEN_FILE, nameFileDict);
        return -1;
    }

    //Files Size
    long long sizeMD5 = countLines(nameFileMD5);
    long long sizeWordsDICT = GetFileSize(nameFileDict);

    //Struct and params
    string *wordsMD5; 
    char* dict = (char*)malloc(sizeWordsDICT);
    string *result; result=new string[sizeMD5];
    std::vector<std::thread> ThreadVector; 
    long long **rangs; 
    rangs = new long long*[numThreads];
    int recovered = 0;

    //Read hashes to crack, read dict and caculate ranges 
    readFile(wordsMD5,nameFileMD5,sizeMD5);
    readFileDict(dict,nameFileDict,sizeWordsDICT);
    fillSplit(rangs,sizeWordsDICT,numThreads);

    cout<<"Running process with "<<numThreads<<" threads..."<<endl;

    auto begin=high_resolution_clock::now();

    
    for(int thread=0; thread<numThreads; thread++){
        int core = thread;
        ThreadVector.emplace_back([&wordsMD5,&dict,&result,&rangs,core,&sizeWordsDICT,&sizeMD5,&recovered, &trace](){crackPasswords(ref(wordsMD5),ref(dict),ref(result),rangs[core], ref(sizeMD5), ref(sizeWordsDICT), core, ref(recovered), ref(trace));});
    }
    for(auto&t :ThreadVector){
        t.join();
    }

    auto end=high_resolution_clock::now();

    if(trace){
        loadingFunction(1);
        cout<<endl;
    }

    int time = duration_cast<milliseconds>(end-begin).count();

    cout<<"Time to execution: "<<time<<" ms."<<endl;
    cout<<"Recovered passwords: "<<recovered<<endl;

    ThreadVector.clear();   
    delete [] dict;
    delete [] wordsMD5;
    for (int i=0; i<numThreads; i++){
        delete[] rangs[i];
    }
    delete[] rangs;

    cout<<"Writing in file the recovered passwords..."<<endl;
    outputFileCrack(result,sizeMD5);
    cout<<"Writing in file the hash passwords in SHA256 with salt..."<<endl;
    outputFileSHA256(result,sizeMD5);

    delete [] result;

    return 0;
}