SNFTools/GBUDBTool/main.cpp

#include <iostream>
#include <string>
#include <sstream>
#include <cstdlib>
#include "../CodeDweller/faults.hpp"
#include "../SNFMulti/GBUdb.hpp"

namespace cd = codedweller;

const std::string VersionInfo = "V0.1";
const int ErrorResultCode = 1;
const int SuccessResultCode = 0;

struct Configuration {
    std::string GBXFile;
    double Probability;
    double Confidence;
    bool findBlack;
    bool showDetails;
    Configuration():
        GBXFile(""),
        Probability(0.0),
        Confidence(0.0),
        findBlack(true),
        showDetails(false) {}
} myConfig;

class BadArguments : public cd::RuntimeFault {public: BadArguments(std::string s):RuntimeFault(s){}};
const BadArguments BadArgumentCount("Wrong number of arguments");
const BadArguments BadSwitch("Bad command line switch");

void setConfiguration(Configuration &conf, int argc, char* argv[]) {
    BadArgumentCount(argc > 6 || argc < 1);
    for(int i = 1; i < argc; i++) {
        std::string thisArgument(argv[i]);

        if(0 == thisArgument.find("-black")) {
            conf.findBlack = true;
        }

        else if(0 == thisArgument.find("-white")) {
            conf.findBlack = false;
        }

        else if(0 == thisArgument.find("-details")) {
            conf.showDetails = true;
        }

        else if(0 == thisArgument.find("-p=")) {
            conf.Probability = strtod(thisArgument.substr(3).c_str(),0);
        }

        else if(0 == thisArgument.find("-c=")) {
            conf.Confidence = strtod(thisArgument.substr(3).c_str(),0);
        }

        else if(0 == thisArgument.find("-")) {
            BadSwitch(true);
        }

        else {
            conf.GBXFile = thisArgument;
        }
    }
}

void displayHelp() {
    std::cerr << "GBUDBTool " << VersionInfo << std::endl;
    std::cerr << "Use: GBUDBTool [-black | -white] [-p=<value>] [-c=<value>] <gbx file>" << std::endl;
    std::cerr << "   -black prints out blacklist IPs based on p & c" << std::endl;
    std::cerr << "   -white prints out whitelist IPs based on p & c" << std::endl;
    std::cerr << "   -details prints out counts, confidence, and probability" << std::endl;
    std::cerr << "   -c=<value> sets the confidence figure for the list" << std::endl;
    std::cerr << "   -p=<value> sets the probability figure for the list" << std::endl;
    std::cerr << "   <gbx file> is the path to a GBUdb snapshot" << std::endl;
}

class Reporter : public GBUdbOperator {
    private:
        Configuration myConfig;

        std::string toIP4String(unsigned int IP) {
            int a, b, c, d;
            d = IP & 0x000000ff; IP >>= 8;
            c = IP & 0x000000ff; IP >>= 8;
            b = IP & 0x000000ff; IP >>= 8;
            a = IP;
            std::ostringstream S;
            S << a << "." << b << "." << c << "." << d;
            return S.str();
        }

    public:

        Reporter(Configuration Config) : myConfig(Config) {}

        GBUdbRecord& operator()(unsigned int IP, GBUdbRecord& R) {
            bool requiredConfidence = (myConfig.Confidence <= R.Confidence());
            bool requiredProbability = false;
            if(myConfig.findBlack) requiredConfidence = (myConfig.Probability <= R.Probability());
            else requiredProbability = (myConfig.Probability >= R.Probability());
            if(requiredProbability && requiredConfidence) {
                std::cout << toIP4String(IP);
                if(myConfig.showDetails) {
                    std::cout << "\t"
                         << "g=" << R.Good()
                         << ", b=" << R.Bad()
                         << ", c=" << R.Confidence()
                         << ", p=" << R.Probability();
                }
                std::cout << std::endl;
            }
            return R;
        }
};

int main(int argc, char* argv[]) {
    try {
        setConfiguration(myConfig, argc, argv);
        GBUdbDataset DB(myConfig.GBXFile.c_str());
        DB.load();
        Reporter R(myConfig);
        DB.doForAllRecords(R);
    }

    catch(const std::exception &e) {
        displayHelp();
        std::cerr << std::endl << "Exception: " << e.what() << std::endl;
        return ErrorResultCode;
    }

    return SuccessResultCode;
}