/*
 * Version: MPL 1.1 / GPLv3+ / LGPLv3+
 *
 * The contents of this file are subject to the Mozilla Public License Version
 * 1.1 (the "License"); you may not use this file except in compliance with
 * the License or as specified alternatively below. You may obtain a copy of
 * the License at http://www.mozilla.org/MPL/
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
 * for the specific language governing rights and limitations under the
 * License.
 *
 * Major Contributor(s):
 * [ Copyright (C) 2011 Tata Consultancy Services, Ltd. Marc-André Laverdière <marc-andre@atc.tcs.com> (initial developer) ]
 *
 * All Rights Reserved.
 *
 * For minor contributions see the git repository.
 *
 * Alternatively, the contents of this file may be used under the terms of
 * either the GNU General Public License Version 3 or later (the "GPLv3+"), or
 * the GNU Lesser General Public License Version 3 or later (the "LGPLv3+"),
 * in which case the provisions of the GPLv3+ or the LGPLv3+ are applicable
 * instead of those above.
 */

#include <unistd.h>
#include <stdio.h>
#include <string>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <boost/format.hpp>
#include <boost/program_options.hpp>
#include <boost/filesystem.hpp>
#include <boost/filesystem/fstream.hpp>

using namespace std;
using namespace boost::program_options;
using namespace boost::filesystem;
using namespace boost::filesystem3;
using namespace boost;

//constants
const string SEEDS_SUBFOLDER = "data";
const string RESULTS_SUBFOLDER = "results";
const string VALGRIND_COMMAND = "valgrind --tool=memcheck";

const string OPTION_HELP = "help";
const string OPTION_PROGRAM = "program";
const string OPTION_DIR = "dir";

const string TERMINATOR = "***^^^FILE_OVER@@@$$$$";

const unsigned char PATTERNS[] = {0xFF, 0xEF};

FILE * startProcess(string& command, string& directory)
{
    stringstream concatenator;
    concatenator << VALGRIND_COMMAND << " " << command << " " << directory;
    return popen(concatenator.str().c_str(), "r");
}

string readProgramOutput(FILE * pipe, string& currentbuffer)
{
    char buffer[128];
    string result = "";
    while(!feof(pipe)) {
        if(fgets(buffer, 128, pipe) != NULL)
        {
            string sBuffer = string(buffer);
            size_t terminatorLocation = sBuffer.find_last_of(TERMINATOR);
            if (terminatorLocation < sBuffer.npos)
            { //we have a match
                //do something
                result += sBuffer.substr(0,terminatorLocation);
                break;
            }
            else
                result += sBuffer;
        }
    }

    return result;
}

bool ensureDirectoryExists(const string& name)
{

    path folder(RESULTS_SUBFOLDER);
    if (!exists(folder))
    {
        create_directory(folder);
    }
    else if (!is_directory(folder))
    {
        cerr << "Error: " <<  name << " is not a directory" << endl;
        return false;
    }

    return true;
}

string createBuggerFileName(string & base, uintmax_t offset, unsigned char pattern)
{
    stringstream concatenator;
    concatenator << base << "-" << dec << offset << "-pattern-" <<
        hex << setiosflags (ios_base::showbase | ios_base::uppercase) << static_cast<unsigned int>(pattern);
    return concatenator.str();
}

string getRandomFileName()
{
    boost::filesystem3::ifstream in("/dev/random");
    stringstream concatenator;
    concatenator <<hex;
    for (int i = 0; i < 10; i++)
        concatenator << in.get();
}

int main(int argc, const char * argv[])
{
    options_description desc("Expected arguments:");
    desc.add_options()
        (OPTION_HELP.c_str(), "Shows the help message")
        (OPTION_PROGRAM.c_str(), value<vector<string> >(), "Program to run under Valgrind")
        (OPTION_DIR.c_str(), value<vector<string> >(), "Directory to contain fuzzed files to test");
    //Read the arguments
    variables_map arguments;
    store(parse_command_line(argc, argv, desc), arguments);


    if (arguments.count(OPTION_PROGRAM) && arguments.count(OPTION_DIR))
    {
        string program = arguments[OPTION_PROGRAM].as<vector<string> >()[0];
        string dir = arguments[OPTION_DIR].as<vector <string> >()[0];
        cout << "Running fuzzer " << "\n" <<
                "\tSeeds directory: " << SEEDS_SUBFOLDER << "\n" <<
                "\tProgram: " << program << "\n" <<
                "\tTemporary directory: " << dir << "\n" <<
                "\tOutput directory: " << RESULTS_SUBFOLDER << endl;

        //Ensure output and temp directories exists. If not, create it.
        if (!ensureDirectoryExists(RESULTS_SUBFOLDER))
            exit(1);
        if (!ensureDirectoryExists(dir))
            exit(1);

        //Load all that stuff
        path seeds_folder(SEEDS_SUBFOLDER);
        if (!exists(seeds_folder))
        {
            cerr << "Error: seeds' folder (" << SEEDS_SUBFOLDER << ") does not exist" << endl;
            return 1;
        }
        else if (!is_directory(seeds_folder))
        {
            cerr << "Error: " << SEEDS_SUBFOLDER << " is not a directory" << endl;
            return 1;
        }

        //Start the fuzzing process
        FILE * fuzzProcess = startProcess(program, dir);
        if (fuzzProcess == NULL || feof(fuzzProcess))
        {
            cerr << "Unable to start sub-process " << endl;
            exit(1);
        }

        //Create directory objects
        path results_folder(RESULTS_SUBFOLDER);
        path tmp_folder(dir);

        //Iterate through all the seeds
        directory_iterator end;
        for(directory_iterator seeds(seeds_folder); seeds != end; seeds++)
        {
            path curr_seed = seeds->path().filename();
            path out_base = results_folder/=curr_seed;
            cout << "Fuzzing " << curr_seed << ". Buggers stored in " << out_base.string() << "-*" <<  endl;

            boost::filesystem3::ifstream in(seeds->path());

            uintmax_t seed_size = file_size(*seeds);

            char * seed_copy = new char[seed_size];
            in >> seed_copy;


            //#1: Fuzz by ORing the patterns
            for (uintmax_t offset = 0; offset < seed_size; offset++)
            {
                for (unsigned char i = 0; i < sizeof(PATTERNS); i++)
                {
                    path samplePath = tmp_folder /= getRandomFileName();
                    boost::filesystem3::ofstream out(samplePath);

                    if (offset > 0){
                        out.write(seed_copy, offset);
                    }
                    out<< (seed_copy[offset] & PATTERNS[i]);
                    if (offset < seed_size -1)
                    {
                        out.write(&seed_copy[offset+1], seed_size-offset);
                    }
                    out.close();
                } //end for all the patterns
            } //end for all the offsets

            in.close();
            delete[] seed_copy;
        } //end for

        //clean up
        fclose (fuzzProcess);


        //#2: Fuzz by 0ing
        //#3: Fuzz by padding

    }
    else
    {
        cout << desc << "\n";
        return 1;
    }
   return 0;
}