Merge branch 'moloch'

Conflicts:
	Untwister.cpp
	prngs/Ruby.cpp

Makefile

@@ -1,13 +1,18 @@
 # Standard flags
-CPPFLAGS = -std=gnu++11 -O3 -pthread -g3 -Wall -c -fmessage-length=0 -MMD -fPIC
+CPPFLAGS = -std=gnu++11 -O3 -g3 -Wall -c -fmessage-length=0 -MMD -fPIC
 PYTHON = /usr/include/python2.7
 BOOST_INC = /usr/include
+OBJS = ./prngs/LSBState.o ./prngs/GlibcRand.o ./prngs/Mt19937.o ./prngs/Ruby.o ./PRNGFactory.o ./Untwister.o
 
+all: GlibcRand Mt19937 RubyRand LSBState PRNGFactory Untwister
+	# Make the cli binary
+	g++ $(CPPFLAGS) -pthread -MF"main.d" -MT"main.d" -o "main.o" "./main.cpp"
+	g++ -std=gnu++11 -O3 -pthread $(OBJS) main.o -o untwister
 
-all: GlibcRand Mt19937 RubyRand LSBState PRNGFactory
-	# Make the binary
-	g++ $(CPPFLAGS) -MF"untwister.d" -MT"untwister.d" -o "untwister.o" "./untwister.cpp"
-	g++ -std=gnu++11 -O3 -pthread -o "untwister" ./prngs/LSBState.o ./prngs/GlibcRand.o ./prngs/Mt19937.o ./prngs/Ruby.o ./PRNGFactory.o ./untwister.o
+python: GlibcRand Mt19937 RubyRand LSBState PRNGFactory Untwister
+	# Make the shared object
+	g++ $(CPPFLAGS) -I$(PYTHON) -I$(BOOST_INC) py-untwister.cpp -o py-untwister.o
+	g++ -std=c++11 -shared -fPIC -O3 $(OBJS) py-untwister.o -lboost_python -lpython2.7 -o untwister.so
 
 GlibcRand:
 	g++ $(CPPFLAGS) -MF"prngs/GlibcRand.d" -MT"prngs/GlibcRand.d" -o "prngs/GlibcRand.o" "./prngs/GlibcRand.cpp"
@@ -23,14 +28,9 @@ LSBState:
 
 PRNGFactory:
 	g++ $(CPPFLAGS) -MF"PRNGFactory.d" -MT"PRNGFactory.d" -o "PRNGFactory.o" "./PRNGFactory.cpp"
-	g++ $(CPPFLAGS) -MF"untwister.d" -MT"untwister.d" -o "untwister.o" "./untwister.cpp"
-	g++ -std=gnu++11 -O3 -pthread -o "untwister" ./prngs/LSBState.o ./prngs/GlibcRand.o ./prngs/Mt19937.o ./prngs/Ruby.o ./PRNGFactory.o ./untwister.o
 
-python: GlibcRand Mt19937 RubyRand LSBState PRNGFactory
-	# Make the shared object
-	g++ $(CPPFLAGS) -I$(PYTHON) -I$(BOOST_INC) py-untwister.cpp -o py-untwister.o
-	g++ -std=c++11 -shared -fPIC -O3 ./prngs/LSBState.o ./prngs/GlibcRand.o ./prngs/Mt19937.o ./prngs/Ruby.o ./PRNGFactory.o py-untwister.o \
-		-lboost_python -lpython2.7 -o untwister.so
+Untwister:
+	g++ $(CPPFLAGS) -pthread -MF"Untwister.d" -MT"Untwister.d" -o "Untwister.o" "./Untwister.cpp"
 
 
 clean:

README.md

@@ -1,13 +1,14 @@
 Untwister
 =========
 
-Multi-threaded seed recovery tool for common PRNGs
+Multi-threaded seed recovery tool for common PRNGs.
+
 
 Supported PRNGs
 =================
-* Glibc rand()
-* Mersenne Twister (mt19937)
-* Ruby's MT-variant rand()
+* Glibc's rand()
+* Mersenne Twister (MT19937)
+* Ruby's MT-variant DEFAULT::rand()
 
 Usage
 ========
@@ -53,9 +54,19 @@ Examples:
 Python Bindings
 =================
 * Python 2.7
-* Requires Boost C++ Python library (use apt-get or brew)
+* Requires Boost C++ Python library
+
+### Ubuntu/Debian Linux
 
 ```
+sudo apt-get install python-dev libboost-python-dev
+make python
+```
+
+### OSX
+
+```
+brew install boost --with-python
 make python
 ```
 

Untwister.cpp

@@ -15,62 +15,106 @@
     along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
 
-#ifndef UNTWISTER_H_
-#define UNTWISTER_H_
-
-#include <stdlib.h>
-#include <limits.h>
-#include <stdint.h>
-#include <iostream>
-#include <fstream>
-#include <string>
-#include <vector>
-#include <thread>
-#include <atomic>
-
-#include "ConsoleColors.h" // TODO: Sepearte logic from stdout calls
-#include "PRNGFactory.h"
-#include "prngs/PRNG.h"
-
-using std::chrono::seconds;
-using std::chrono::milliseconds;
-using std::chrono::duration_cast;
-using std::chrono::steady_clock;
-
-// Pair of <seed, quality of fit>
-typedef std::pair<uint32_t, double> Seed;
-
-static std::vector<uint32_t> observedOutputs;
-static const unsigned int ONE_YEAR = 31536000;
+#include "Untwister.h"
 
 
-/* Yeah lots of parameters, but such is the life of a thread */
-void BruteForce(unsigned int id, std::atomic<bool> *isCompleted, std::vector<std::vector<Seed>* > *answers,
-        std::vector<uint32_t> *status, double minimumConfidence, uint32_t startingSeed, uint32_t endingSeed,
-        uint32_t depth, std::string rng)
+Untwister::Untwister()
 {
+    isCompleted = new std::atomic<bool>(false);
+    isRunning = new std::atomic<bool>(false);
+    observedOutputs = new std::vector<uint32_t>;
+    depth = DEFAULT_DEPTH;
+    minConfidence = DEFAULT_MIN_CONFIDENCE;
+    threads = std::thread::hardware_concurrency();
+    prng = "glibc-rand";
+    answers = new std::vector<std::vector<Seed>* >(threads);
+    status = new std::vector<uint32_t>(threads);
+}
+
+Untwister::Untwister(unsigned int observationSize)
+{
+    isCompleted = new std::atomic<bool>(false);
+    isRunning = new std::atomic<bool>(false);
+    observedOutputs = new std::vector<uint32_t>(observationSize);
+    depth = DEFAULT_DEPTH;
+    minConfidence = DEFAULT_MIN_CONFIDENCE;
+    threads = std::thread::hardware_concurrency();
+    prng = "glibc-rand";
+    answers = new std::vector<std::vector<Seed>* >(threads);
+    status = new std::vector<uint32_t>(threads);
+}
+
+Untwister::~Untwister()
+{
+    delete isCompleted;
+    delete isRunning;
+    delete observedOutputs;
+    delete answers;
+    delete status;
+}
+
+std::vector<Seed> Untwister::bruteforce(uint32_t lowerBoundSeed, uint32_t upperBoundSeed)
+{
+
+    std::vector<uint32_t> labor = divisionOfLabor(upperBoundSeed - lowerBoundSeed, threads);
+    uint32_t startAt = lowerBoundSeed;
+
+    std::vector<std::thread> pool = std::vector<std::thread>(threads);
+    for(unsigned int id = 0; id < threads; ++id)
+    {
+        uint32_t endAt = startAt + labor.at(id);
+        pool[id] = std::thread(&Untwister::worker, this, id, startAt, endAt);
+        startAt += labor.at(id);
+    }
+
+    isRunning->store(true, std::memory_order_relaxed);
+
+    for(unsigned int id = 0; id < pool.size(); ++id)
+    {
+        pool[id].join();
+    }
+
+    std::vector<Seed> results = std::vector<Seed>();
+    for (unsigned int id = 0; id < answers->size(); ++id)
+    {
+        for (unsigned int index = 0; index < answers->at(id)->size(); ++index)
+        {
+            results.push_back(answers->at(id)->at(index));
+        }
+        delete answers->at(id);
+    }
+
+    return results;
+}
+
+/* This method is execute as a seperate thread */
+void Untwister::worker(unsigned int id, uint32_t startingSeed, uint32_t endingSeed)
+{
+
     /* Each thread must have a local factory unless you like mutexes and/or segfaults */
     PRNGFactory factory;
-    PRNG *generator = factory.getInstance(rng);
-    answers->at(id) = new std::vector<Seed>;
+    PRNG *generator = factory.getInstance(prng);
+    answers->at(id) = new std::vector<Seed>();
 
     for(uint32_t seedIndex = startingSeed; seedIndex <= endingSeed; ++seedIndex)
     {
 
         if(isCompleted->load(std::memory_order_relaxed))
+        {
             break;  // Some other thread found the seed
+        }
 
         generator->seed(seedIndex);
         uint32_t matchesFound = 0;
         for(uint32_t index = 0; index < depth; index++)
         {
             uint32_t nextRand = generator->random();
-            uint32_t observed = observedOutputs[matchesFound];
+            uint32_t observed = observedOutputs->at(matchesFound);
 
             if(observed == nextRand)
             {
                 matchesFound++;
-                if(matchesFound == observedOutputs.size())
+                if(matchesFound == observedOutputs->size())
                 {
                     break;  // This seed is a winner if we get to the end
                 }
@@ -78,13 +122,13 @@ void BruteForce(unsigned int id, std::atomic<bool> *isCompleted, std::vector<std
         }
 
         status->at(id) = seedIndex - startingSeed;
-        double confidence = ((double) matchesFound / (double) observedOutputs.size()) * 100.0;
-        if(minimumConfidence <= confidence)
+        double confidence = ((double) matchesFound / (double) observedOutputs->size()) * 100.0;
+        if(minConfidence <= confidence)
         {
             Seed seed = Seed(seedIndex, confidence);
             answers->at(id)->push_back(seed);
         }
-        if(matchesFound == observedOutputs.size())
+        if(matchesFound == observedOutputs->size())
         {
             isCompleted->store(true, std::memory_order_relaxed);
         }
@@ -92,51 +136,6 @@ void BruteForce(unsigned int id, std::atomic<bool> *isCompleted, std::vector<std
     delete generator;
 }
 
-/* Divide X number of seeds among Y number of threads */
-std::vector<uint32_t> DivisionOfLabor(uint32_t sizeOfWork, uint32_t numberOfWorkers)
-{
-    uint32_t work = sizeOfWork / numberOfWorkers;
-    uint32_t leftover = sizeOfWork % numberOfWorkers;
-    std::vector<uint32_t> labor(numberOfWorkers);
-    for(uint32_t index = 0; index < numberOfWorkers; ++index)
-    {
-        if(0 < leftover)
-        {
-            labor[index] = work + 1;
-            --leftover;
-        }
-        else
-        {
-            labor[index] = work;
-        }
-    }
-    return labor;
-}
-
-/* Generic Threading */
-void StartBruteForce(unsigned int threads, std::vector<std::vector<Seed>* > *answers, double minimumConfidence,
-        uint32_t lowerBoundSeed, uint32_t upperBoundSeed, uint32_t depth, std::string rng)
-{
-    std::atomic<bool> *isCompleted = new std::atomic<bool>(false);
-    std::vector<std::thread> pool(threads);
-    std::vector<uint32_t> *status = new std::vector<uint32_t>(threads);
-    std::vector<uint32_t> labor = DivisionOfLabor(upperBoundSeed - lowerBoundSeed, threads);
-    uint32_t startAt = lowerBoundSeed;
-
-    for(unsigned int id = 0; id < threads; ++id)
-    {
-        uint32_t endAt = startAt + labor.at(id);
-        pool[id] = std::thread(BruteForce, id, isCompleted, answers, status, minimumConfidence, startAt, endAt, depth, rng);
-        startAt += labor.at(id);
-    }
-
-    for(unsigned int id = 0; id < pool.size(); ++id)
-    {
-        pool[id].join();
-    }
-    delete status;
-}
-
 
 /*
     This is the "smarter" method of breaking RNGs. We use consecutive integers
@@ -144,18 +143,17 @@ void StartBruteForce(unsigned int threads, std::vector<std::vector<Seed>* > *ans
     method, however, we won't typically recover an actual seed value.
     But the effect is the same.
 */
-bool InferState(const std::string& rng)
+bool Untwister::inferState()
 {
-    std::cout << INFO << "Trying state inference" << std::endl;
 
     PRNGFactory factory;
-    PRNG *generator = factory.getInstance(rng);
+    PRNG *generator = factory.getInstance(prng);
     uint32_t stateSize = generator->getStateSize();
 
-    if(observedOutputs.size() <= stateSize)
+    if(observedOutputs->size() <= stateSize)
     {
-        std::cout << WARN << "Not enough observed values to perform state inference." << std::endl;
-        std::cout << WARN << "Try again with more than " << stateSize << " values" << std::endl;
+        std::cout << WARN << "Not enough observed values to perform state inference,"
+                  << " try again with more than " << stateSize << " values." << std::endl;
         return false;
     }
 
@@ -164,36 +162,36 @@ bool InferState(const std::string& rng)
     /* Guaranteed from the above to loop at least one time */
     std::vector<double> scores;
     std::vector<uint32_t> best_state;
-    for(uint32_t i = 0; i < (observedOutputs.size() - stateSize); i++)
+    for(uint32_t index = 0; index < (observedOutputs->size() - stateSize); ++index)
     {
-        std::vector<uint32_t>::const_iterator first = observedOutputs.begin() + i;
-        std::vector<uint32_t>::const_iterator last = observedOutputs.begin() + i + stateSize;
+        std::vector<uint32_t>::const_iterator first = observedOutputs->begin() + index;
+        std::vector<uint32_t>::const_iterator last = observedOutputs->begin() + index + stateSize;
         std::vector<uint32_t> state(first, last);
 
         /* Make predictions based on the state */
         std::vector<uint32_t> evidenceForward
-            ((std::vector<uint32_t>::const_iterator)observedOutputs.begin(), first);
+            ((std::vector<uint32_t>::const_iterator)observedOutputs->begin(), first);
         std::vector<uint32_t> evidenceBackward
-            (last+1, (std::vector<uint32_t>::const_iterator)observedOutputs.end());
+            (last+1, (std::vector<uint32_t>::const_iterator)observedOutputs->end());
         generator->setState(state);
 
         /* Provide additional evidence for tuning on PRNGs that require it */
-        generator->setEvidence(observedOutputs);
+        generator->setEvidence((*observedOutputs));
         generator->tune(evidenceForward, evidenceBackward);
 
         std::vector<uint32_t> predictions_forward =
-            generator->predictForward(((observedOutputs.size() - stateSize) - i));
-        std::vector<uint32_t> predictions_backward =
-            generator->predictBackward(i);
+            generator->predictForward(((observedOutputs->size() - stateSize) - index));
+
+        std::vector<uint32_t> predictions_backward = generator->predictBackward(index);
 
         /* Test the prediction against the rest of the observed data */
         /* Forward */
         uint32_t matchesFound = 0;
         uint32_t index_pred = 0;
-        uint32_t index_obs = i + stateSize;
-        while(index_obs < observedOutputs.size() && index_pred < predictions_forward.size())
+        uint32_t index_obs = index + stateSize;
+        while(index_obs < observedOutputs->size() && index_pred < predictions_forward.size())
         {
-            if(observedOutputs[index_obs] == predictions_forward[index_pred])
+            if(observedOutputs->at(index_obs) == predictions_forward[index_pred])
             {
                 matchesFound++;
                 index_obs++;
@@ -203,10 +201,10 @@ bool InferState(const std::string& rng)
 
         /* Backward */
         index_pred = 0;
-        index_obs = i;
+        index_obs = index;
         while(index_obs > 0 && index_pred < predictions_backward.size())
         {
-            if(observedOutputs[index_obs] == predictions_backward[index_pred])
+            if(observedOutputs->at(index_obs) == predictions_backward[index_pred])
             {
                 matchesFound++;
                 index_obs--;
@@ -215,7 +213,7 @@ bool InferState(const std::string& rng)
         }
 
         /* If we get a perfect guess, then try reversing out the seed, and exit */
-        if(matchesFound == (observedOutputs.size() - stateSize))
+        if(matchesFound == (observedOutputs->size() - stateSize))
         {
             uint32_t outSeed = 0;
             if(generator->reverseToSeed(&outSeed, 10000))
@@ -235,7 +233,7 @@ bool InferState(const std::string& rng)
             return true;
         }
 
-        double score = (double)(matchesFound*100) / (double)(observedOutputs.size() - stateSize);
+        double score = (double)(matchesFound*100) / (double)(observedOutputs->size() - stateSize);
         scores.push_back(score);
         if(score > highscore)
         {
@@ -259,33 +257,42 @@ bool InferState(const std::string& rng)
         std::cout << INFO << "State Inference failed" << std::endl;
     }
 
+    delete generator;
     return false;
 }
 
-/* For easier testing, will generate a series of random numbers at a given seed and depth*/
-std::vector<uint32_t> GenerateSample(uint32_t seed, uint32_t depth, std::string rng)
+std::vector<uint32_t> Untwister::generateSampleFromSeed(uint32_t seed)
 {
     PRNGFactory factory;
-    PRNG *generator = factory.getInstance(rng);
+    PRNG *generator = factory.getInstance(prng);
     generator->seed(seed);
+    PRNG *distance_gen = factory.getInstance(prng);
+    distance_gen->seed(time(NULL));
+    uint32_t distance = distance_gen->random() % (depth - 10);
 
-    std::vector<uint32_t> results;
-    for (unsigned int index = 0; index < depth; ++index)
+    // Burn a bunch of random numbers
+    for (uint32_t index = 0; index < distance; ++index)
+    {
+        generator->random();
+    }
+
+    std::vector<uint32_t> results(10);
+    for (unsigned int index = 0; index < 10; ++index)
     {
         results.push_back(generator->random());
     }
     delete generator;
+    delete distance_gen;
     return results;
 }
 
-/* For easier testing, will generate a series of random numbers at a given state and depth */
-std::vector<uint32_t> GenerateSample(std::vector<uint32_t> state, uint32_t depth, std::string rng)
+std::vector<uint32_t> Untwister::generateSampleFromState()
 {
     PRNGFactory factory;
-    PRNG *generator = factory.getInstance(rng);
-    generator->setState(state);
+    PRNG *generator = factory.getInstance(prng);
+    generator->setState((*observedOutputs));
 
-    std::vector<uint32_t> results;
+    std::vector<uint32_t> results(depth);
     for (unsigned int index = 0; index < depth; ++index)
     {
         results.push_back(generator->random());
@@ -294,4 +301,121 @@ std::vector<uint32_t> GenerateSample(std::vector<uint32_t> state, uint32_t depth
     return results;
 }
 
-#endif /* UNTWISTER_H_ */
+/* Divide X work among Y number of threads, and evenly distribute remainders */
+std::vector<uint32_t> Untwister::divisionOfLabor(uint32_t sizeOfWork, uint32_t numberOfWorkers)
+{
+    uint32_t work = sizeOfWork / numberOfWorkers;
+    uint32_t leftover = sizeOfWork % numberOfWorkers;
+    std::vector<uint32_t> labor(numberOfWorkers);
+    for(uint32_t index = 0; index < numberOfWorkers; ++index)
+    {
+        if (0 < leftover)
+        {
+            labor[index] = work + 1;
+            --leftover;
+        }
+        else
+        {
+            labor[index] = work;
+        }
+    }
+    return labor;
+}
+
+void Untwister::addObservedOutput(uint32_t observedOutput)
+{
+    observedOutputs->push_back(observedOutput);
+}
+
+std::vector<uint32_t>* Untwister::getObservedOutputs()
+{
+    return observedOutputs;
+}
+
+std::vector<uint32_t>* Untwister::getStatus()
+{
+    return status;
+}
+
+std::vector<std::string> Untwister::getPRNGNames()
+{
+    PRNGFactory factory;
+    return factory.getNames();
+}
+
+
+void Untwister::setPRNG(std::string prng)
+{
+    if (isSupportedPRNG(prng))
+    {
+        this->prng = prng;
+    }
+}
+
+void Untwister::setPRNG(char *prng)
+{
+    setPRNG(std::string(prng));
+}
+
+std::string Untwister::getPRNG()
+{
+    return prng;
+}
+
+bool Untwister::isSupportedPRNG(char *prng)
+{
+    return isSupportedPRNG(std::string(prng));
+}
+
+bool Untwister::isSupportedPRNG(std::string prng)
+{
+    std::vector<std::string> names = getPRNGNames();
+    return std::find(names.begin(), names.end(), prng) == names.end() ? false:true;
+}
+
+void Untwister::setThreads(unsigned int threads)
+{
+    if (!isRunning->load(std::memory_order_relaxed))
+    {
+        this->threads = threads;
+        delete answers;
+        delete status;
+        answers = new std::vector<std::vector<Seed>* >(threads);
+        status = new std::vector<uint32_t>(threads);
+    }
+}
+
+unsigned int Untwister::getThreads()
+{
+    return threads;
+}
+
+void Untwister::setDepth(unsigned int depth)
+{
+    this->depth = depth;
+}
+
+unsigned int Untwister::getDepth()
+{
+    return depth;
+}
+
+void Untwister::setMinConfidence(double minConfidence)
+{
+    this->minConfidence = minConfidence;
+}
+
+double Untwister::getMinConfidence()
+{
+    return minConfidence;
+}
+
+std::atomic<bool>* Untwister::getIsCompleted()
+{
+    return isCompleted;
+}
+
+std::atomic<bool>* Untwister::getIsRunning()
+{
+    return isRunning;
+}

Untwister.h

@@ -0,0 +1,91 @@
+/*
+    Copyright Dan Petro, 2014
+
+    This program 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.
+
+    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 General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+#ifndef UNTWISTER_H_
+#define UNTWISTER_H_
+
+#include <climits>
+#include <cstdlib>
+#include <cstdint>
+#include <string>
+#include <vector>
+#include <thread>
+#include <atomic>
+#include <fstream>
+#include <iostream>
+
+#include "ConsoleColors.h"
+#include "PRNGFactory.h"
+#include "prngs/PRNG.h"
+
+// Pair of <seed, quality of fit>
+typedef std::pair<uint32_t, double> Seed;
+
+static const uint32_t DEFAULT_DEPTH = 1000;
+static const double DEFAULT_MIN_CONFIDENCE = 100.0;
+
+class Untwister
+{
+
+public:
+    Untwister();
+    Untwister(unsigned int threads);
+    Untwister(unsigned int threads, unsigned int observationSize);
+    virtual ~Untwister();
+
+    std::vector<Seed> bruteforce(uint32_t lowerBoundSeed, uint32_t upperBoundSeed);
+    bool inferState();
+
+    std::vector<std::string> getPRNGNames();
+    void setPRNG(std::string prng);
+    void setPRNG(char *prng);
+    std::string getPRNG();
+    bool isSupportedPRNG(std::string prng);
+    bool isSupportedPRNG(char* prng);
+
+    void setMinConfidence(double minConfidence);
+    double getMinConfidence();
+    void setDepth(uint32_t depth);
+    uint32_t getDepth();
+    void setThreads(unsigned int threads);
+    unsigned int getThreads();
+    void addObservedOutput(uint32_t observedOutput);
+    std::vector<uint32_t>* getObservedOutputs();
+    std::vector<uint32_t>* getStatus();
+    std::atomic<bool>* getIsCompleted();
+    std::atomic<bool>* getIsRunning();
+
+    std::vector<uint32_t> generateSampleFromSeed(uint32_t seed);
+    std::vector<uint32_t> generateSampleFromState();
+
+private:
+    unsigned int threads;
+    double minConfidence;
+    uint32_t depth;
+    std::string prng;
+    std::atomic<bool> *isCompleted;
+    std::atomic<bool> *isRunning;
+    std::vector<uint32_t> *status;
+    std::vector<std::vector<Seed>* > *answers;
+    std::vector<uint32_t> *observedOutputs;
+
+    void worker(unsigned int id, uint32_t startingSeed, uint32_t endingSeed);
+    std::vector<uint32_t> divisionOfLabor(uint32_t sizeOfWork, uint32_t numberOfWorkers);
+
+};
+
+#endif /* UNTWISTER_H_ */

main.cpp

@@ -18,12 +18,34 @@
 
 #include <stdio.h>
 #include <getopt.h>
+#include <execinfo.h>
+#include <signal.h>
+#include <chrono>
+#include <atomic>
+#include <thread>
 
-#include "untwister.h"
 #include "ConsoleColors.h"
+#include "Untwister.h"
 
+using std::chrono::seconds;
+using std::chrono::milliseconds;
+using std::chrono::duration;
+using std::chrono::duration_cast;
+using std::chrono::steady_clock;
 
-void Usage(PRNGFactory factory, unsigned int threads)
+static const unsigned int ONE_YEAR = 31536000;
+static const unsigned int TRACE_SIZE = 10;
+
+/* Segfault handler - for debugging only */
+void handler(int sig) {
+    void *trace[TRACE_SIZE];
+    size_t size = backtrace(trace, TRACE_SIZE);
+    std::cerr << "[!] SIGSEGV: " << sig << std::endl;
+    backtrace_symbols_fd(trace, size, 2);
+    exit(1);
+}
+
+void Usage(Untwister *untwister)
 {
     std::cout << BOLD << "Untwister" << RESET << " - Recover PRNG seeds from observed values." << std::endl;
     std::cout << "\t-i <input_file> [-d <depth> ] [-r <prng>] [-g <seed>] [-t <threads>] [-c <confidence>]\n" << std::endl;
@@ -34,85 +56,115 @@ void Usage(PRNGFactory factory, unsigned int threads)
     std::cout << "\t\tChoosing a higher depth value will make brute forcing take longer (linearly), but is" << std::endl;
     std::cout << "\t\trequired for cases where the generator has been used many times already." << std::endl;
     std::cout << "\t-r <prng>\n\t\tThe PRNG algorithm to use. Supported PRNG algorithms:" << std::endl;
-    std::vector<std::string> names = factory.getNames();
+    std::vector<std::string> names = untwister->getPRNGNames();
     for (unsigned int index = 0; index < names.size(); ++index)
     {
         std::cout << "\t\t" << BOLD << " * " << RESET << names[index];
         if (index == 0)
+        {
             std::cout << " (default)";
+        }
         std::cout << std::endl;
     }
     std::cout << "\t-u\n\t\tUse bruteforce, but only for unix timestamp values within a range of +/- 1 " << std::endl;
     std::cout << "\t\tyear from the current time." << std::endl;
     std::cout << "\t-g <seed>\n\t\tGenerate a test set of random numbers from the given seed (at a random depth)" << std::endl;
     std::cout << "\t-c <confidence>\n\t\tSet the minimum confidence percentage to report" << std::endl;
-    std::cout << "\t-t <threads>\n\t\tSpawn this many threads (default is " << threads << ")" << std::endl;
+    std::cout << "\t-t <threads>\n\t\tSpawn this many threads (default is " << untwister->getThreads() << ")" << std::endl;
     std::cout << std::endl;
 }
 
-void StatusThread(std::vector<std::thread>& pool, bool& isCompleted, uint32_t totalWork, std::vector<uint32_t> *status)
+
+void DisplayProgress(Untwister *untwister, uint32_t totalWork)
 {
-    double percent = 0;
-    steady_clock::time_point start = steady_clock::now();
-    while (!isCompleted)
+    std::atomic<bool>* isRunning = untwister->getIsRunning();
+    while (!isRunning->load(std::memory_order_relaxed))
+    {
+        std::this_thread::sleep_for(milliseconds(100));
+    }
+
+    double percent = 0.0;
+    double seedsPerSec = 0.0;
+    double timeLeft = 0.0;
+    steady_clock::time_point started = steady_clock::now();
+    std::vector<uint32_t> *status = untwister->getStatus();
+    std::atomic<bool> *isCompleted = untwister->getIsCompleted();
+    char spinner[] = {'|', '/', '-', '\\'};
+    unsigned int count = 0;
+    while (!isCompleted->load(std::memory_order_relaxed))
     {
         unsigned int sum = 0;
+        duration<double> time_span = duration_cast<duration<double>>(steady_clock::now() - started);
         for (unsigned int index = 0; index < status->size(); ++index)
         {
             sum += status->at(index);
         }
         percent = ((double) sum / (double) totalWork) * 100.0;
-        isCompleted = (100.0 <= percent);
-        std::cout << CLEAR << DEBUG << "Progress: " << percent << '%';
-        std::cout << " (" << (int) duration_cast<seconds>(steady_clock::now() - start).count() << " seconds)";
+        if (0 < time_span.count())
+        {
+            seedsPerSec = (double) sum / (double) time_span.count();
+            if (0 == count % 20)
+            {
+                timeLeft = ((double) (totalWork - sum) / seedsPerSec) / 60.0;
+            }
+        }
+
+        std::cout << CLEAR << BOLD << PURPLE << "[" << spinner[count % 4] << "]" << RESET
+                  << " Progress: " << percent << '%'
+                  << "  [" << sum << " / " << totalWork << "]"
+                  << "  ~" << seedsPerSec << "/sec"
+                  << "  " << timeLeft << " minute(s)";
         std::cout.flush();
-        std::this_thread::sleep_for(milliseconds(150));
+        ++count;
+        std::this_thread::sleep_for(milliseconds(100));
     }
     std::cout << CLEAR;
 }
 
-void FindSeed(const std::string& rng, unsigned int threads, double minimumConfidence, uint32_t lowerBoundSeed,
-        uint32_t upperBoundSeed, uint32_t depth)
+void FindSeed(Untwister *untwister, uint32_t lowerBoundSeed, uint32_t upperBoundSeed)
 {
-    std::cout << INFO << "Looking for seed using " << rng << std::endl;
-    std::cout << INFO << "Spawning " << threads << " worker thread(s) ..." << std::endl;
+    std::cout << INFO << "Looking for seed using " << BOLD << untwister->getPRNG() << RESET << std::endl;
+    std::cout << INFO << "Spawning " << untwister->getThreads() << " worker thread(s) ..." << std::endl;
 
-
-    /* Each thread needs their own set of answers to avoid locking */
-    std::vector<std::vector<Seed>* > *answers = new std::vector<std::vector<Seed>* >(threads);
     steady_clock::time_point elapsed = steady_clock::now();
-    StartBruteForce(threads, answers, minimumConfidence, lowerBoundSeed, upperBoundSeed, depth, rng);
+    std::thread progressThread(DisplayProgress, untwister, upperBoundSeed - lowerBoundSeed);
 
-    std::cout << INFO << "Completed in " << duration_cast<seconds>(steady_clock::now() - elapsed).count()
+    auto results = untwister->bruteforce(lowerBoundSeed, upperBoundSeed);
+    auto isCompleted = untwister->getIsCompleted();
+    if (!isCompleted->load(std::memory_order_relaxed))
+    {
+        isCompleted->store(true, std::memory_order_relaxed);
+    }
+    progressThread.join();
+
+    /* Total time elapsed */
+    std::cout << INFO << "Completed in "
+              << duration_cast<seconds>(steady_clock::now() - elapsed).count()
               << " second(s)" << std::endl;
 
     /* Display results */
-    for (unsigned int id = 0; id < answers->size(); ++id)
+    for (unsigned int index = 0; index < results.size(); ++index)
     {
-        /* Look for answers from each thread */
-        for (unsigned int index = 0; index < answers->at(id)->size(); ++index)
-        {
-            std::cout << SUCCESS << "Found seed " << answers->at(id)->at(index).first
-                      << " with a confidence of " << answers->at(id)->at(index).second
-                      << '%' << std::endl;
-        }
-        delete answers->at(id);
+        std::cout << SUCCESS << "Found seed " << results[index].first
+                  << " with a confidence of " << results[index].second
+                  << '%' << std::endl;
     }
-    delete answers;
 }
 
 int main(int argc, char *argv[])
 {
+    /* Signal Handlers */
+    signal(SIGSEGV, handler);
+    signal(SIGILL, handler);
+    signal(SIGABRT, handler);
+
     int c;
-    unsigned int threads = std::thread::hardware_concurrency();
+
     uint32_t lowerBoundSeed = 0;
     uint32_t upperBoundSeed = UINT_MAX;
-    uint32_t depth = 1000;
     uint32_t seed = 0;
     bool generateFlag = false;
-    double minimumConfidence = 100.0;
-    PRNGFactory factory;
-    std::string rng = factory.getNames()[0];
+    Untwister *untwister = new Untwister();
 
     while ((c = getopt(argc, argv, "d:i:g:t:r:c:uh")) != -1)
     {
@@ -132,23 +184,30 @@ int main(int argc, char *argv[])
             }
             case 'r':
             {
-                rng = optarg;
-                std::vector<std::string> names = factory.getNames();
-                if (std::find(names.begin(), names.end(), rng) == names.end())
+                if (!untwister->isSupportedPRNG(optarg))
                 {
                     std::cerr << WARN << "ERROR: The PRNG \"" << optarg << "\" is not supported, see -h" << std::endl;
                     return EXIT_FAILURE;
                 }
+                else
+                {
+                    untwister->setPRNG(optarg);
+                }
                 break;
             }
             case 'd':
             {
-                depth = strtoul(optarg, NULL, 10);
+                unsigned int depth = strtoul(optarg, NULL, 10);
                 if (depth == 0)
                 {
                     std::cerr << WARN << "ERROR: Please enter a valid depth > 1" << std::endl;
                     return EXIT_FAILURE;
                 }
+                else
+                {
+                    std::cout << INFO << "Depth set to: " << depth << std::endl;
+                    untwister->setDepth(depth);
+                }
                 break;
             }
             case 'i':
@@ -161,33 +220,43 @@ int main(int argc, char *argv[])
                 std::string line;
                 while (std::getline(infile, line))
                 {
-                    observedOutputs.push_back(strtoul(line.c_str(), NULL, 0));
+                    uint32_t value = strtoul(line.c_str(), NULL, 0);
+                    untwister->addObservedOutput(value);
                 }
                 break;
             }
             case 't':
             {
-                threads = strtoul(optarg, NULL, 10);
+                unsigned int threads = strtoul(optarg, NULL, 10);
                 if (threads == 0)
                 {
                     std::cerr << WARN << "ERROR: Please enter a valid number of threads > 1" << std::endl;
                     return EXIT_FAILURE;
                 }
+                else
+                {
+                    untwister->setThreads(threads);
+                }
                 break;
             }
             case 'c':
             {
-                minimumConfidence = ::atof(optarg);
+                double minimumConfidence = ::atof(optarg);
                 if (minimumConfidence <= 0 || 100.0 < minimumConfidence)
                 {
                     std::cerr << WARN << "ERROR: Invalid confidence percentage " << std::endl;
                     return EXIT_FAILURE;
                 }
+                else
+                {
+                    std::cout << INFO << "Minimum confidence set to: " << minimumConfidence << std::endl;
+                    untwister->setMinConfidence(minimumConfidence);
+                }
                 break;
             }
             case 'h':
             {
-                Usage(factory, threads);
+                Usage(untwister);
                 return EXIT_SUCCESS;
             }
             case '?':
@@ -198,46 +267,51 @@ int main(int argc, char *argv[])
                    std::cerr << "Unknown option `-" << optopt << "'." << std::endl;
                 else
                    std::cerr << "Unknown option character `" << optopt << "'." << std::endl;
-                Usage(factory, threads);
+                Usage(untwister);
                 return EXIT_FAILURE;
             }
             default:
             {
-                Usage(factory, threads);
+                Usage(untwister);
                 return EXIT_FAILURE;
             }
         }
     }
 
-    if(generateFlag)
+    if (generateFlag)
     {
         std::vector<uint32_t> results;
-        if(observedOutputs.empty())
+        if(untwister->getObservedOutputs()->empty())
         {
-            results = GenerateSample(seed, depth, rng);
+            results = untwister->generateSampleFromSeed(seed);
         }
         else
         {
-            results = GenerateSample(observedOutputs, depth, rng);
+            results = untwister->generateSampleFromState();
         }
         for (unsigned int index = 0; index < results.size(); ++index)
+        {
             std::cout << results.at(index) << std::endl;
+        }
         return EXIT_SUCCESS;
 
     }
 
-    if(observedOutputs.empty())
+    if (untwister->getObservedOutputs()->empty())
     {
-        Usage(factory, threads);
+        Usage(untwister);
         std::cerr << WARN << "ERROR: No input numbers provided. Use -i <file> to provide a file" << std::endl;
         return EXIT_FAILURE;
     }
 
-    if(InferState(rng))
+    if (untwister->inferState())
     {
         return EXIT_SUCCESS;
     }
-    FindSeed(rng, threads, minimumConfidence, lowerBoundSeed, upperBoundSeed, depth);
+
+    FindSeed(untwister, lowerBoundSeed, upperBoundSeed);
+    delete untwister;
+
     return EXIT_SUCCESS;
 }
 

py-untwister.cpp

@@ -10,15 +10,17 @@
 #include <boost/python.hpp>
 #include <execinfo.h>
 #include <signal.h>
+#include <iostream>
 
-#include "untwister.h"
+#include "Untwister.h"
 
 using namespace boost::python;
 
 static const unsigned int TRACE_SIZE = 10;
 
 /* Segfault handler - for debugging only */
-void handler(int sig) {
+void handler(int sig)
+{
     void *trace[TRACE_SIZE];
     size_t size = backtrace(trace, TRACE_SIZE);
     std::cerr << "[!] SIGSEGV: " << sig << std::endl;
@@ -30,6 +32,7 @@ void handler(int sig) {
 void PythonInit()
 {
     signal(SIGSEGV, handler);
+    signal(SIGABRT, handler);
     if(!Py_IsInitialized())
     {
         Py_Initialize();
@@ -37,82 +40,64 @@ void PythonInit()
     }
 }
 
-/* Raises a Python exception if the string is not a supported prng */
-void CheckPRNG(std::string prng)
+/* Find a seed */
+list FindSeed(std::string prng, list inputs, unsigned int threads, float minimumConfidence,
+    uint32_t lowerBoundSeed, uint32_t upperBoundSeed, uint32_t depth)
 {
-    PRNGFactory factory;
-    std::vector<std::string> names = factory.getNames();
-    if(std::find(names.begin(), names.end(), prng) == names.end())
+    Untwister *untwister = new Untwister(len(inputs));
+    if(!untwister->isSupportedPRNG(prng))
     {
         /* Raise Python Exception */
         PyErr_SetString(PyExc_ValueError, "Unsupported PRNG");
         throw error_already_set();
     }
-}
-
-/* Find a seed */
-list FindSeed(const std::string& rng, list inputs, unsigned int threads, float minimumConfidence,
-    uint32_t lowerBoundSeed, uint32_t upperBoundSeed, uint32_t depth)
-{
-    CheckPRNG(rng);
+    untwister->setPRNG(prng);
+    untwister->setThreads(threads);
+    untwister->setDepth(depth);
+    untwister->setMinConfidence(minimumConfidence);
 
     /* Convert Python list object to observedOutputs's std::vector<uint32_t> */
-    for(unsigned int index = 0; index < len(inputs); ++index) {
-        unsigned int data = boost::python::extract<unsigned int>(inputs[index]);
-        observedOutputs.push_back(data);
+    for(unsigned int index = 0; index < len(inputs); ++index)
+    {
+        uint32_t data = boost::python::extract<unsigned int>(inputs[index]);
+        untwister->getObservedOutputs()->at(index) = data;
     }
 
-    std::vector<std::vector<Seed>* > *answers = new std::vector<std::vector<Seed>* >(threads);
-
     /* Suspend Python's thread, so we can use native C++ threads */
-    PyThreadState* pyThreadState = PyEval_SaveThread();
-    StartBruteForce(threads, answers, (double) minimumConfidence, lowerBoundSeed, upperBoundSeed, depth, rng);
+    PyThreadState *pyThreadState = PyEval_SaveThread();
+
+    std::vector<std::pair<uint32_t, double> > results = untwister->bruteforce(lowerBoundSeed, upperBoundSeed);
+
 
     /* Clean up and restore Python thread state */
     PyEval_RestoreThread(pyThreadState);
-    pyThreadState = NULL;
 
     /* Covert answers to python list of tuples */
-    list results;
-    for(unsigned int id = 0; id < answers->size(); ++id)
-    {
-        for(unsigned int index = 0; index < answers->at(id)->size(); ++index)
-        {
-            tuple result = make_tuple(answers->at(id)->at(index).first, answers->at(id)->at(index).second);
-            results.append(result);
-        }
-        delete answers->at(id);
-    }
-    delete answers;
-    return results;
-}
-
-/* Generate a sample sequence */
-list Sample(std::string prng, uint32_t seed, uint32_t depth)
-{
-    CheckPRNG(prng);
-
-    std::vector<uint32_t> results = GenerateSample(seed, depth, prng);
-    /* Convert results to a Python object */
-    list sample;
+    list pyResults;
     for(unsigned int index = 0; index < results.size(); ++index)
     {
-        sample.append(results.at(index));
+        tuple pySeed = make_tuple(results[index].first, results[index].second);
+        pyResults.append(pySeed);
     }
-    return sample;
+    delete untwister;
+    return pyResults;
 }
 
 /* List all supported PRNGs */
 list Prngs()
 {
-    PRNGFactory factory;
-    std::vector<std::string> names = factory.getNames();
-    list prngs;
+    Untwister *untwister = new Untwister();
+    std::vector<std::string> names = untwister->getPRNGNames();
+
+    list pyPRNGs;
     for(unsigned int index = 0; index < names.size(); ++index)
     {
-        prngs.append(names.at(index));
+        pyPRNGs.append(names[index]);
     }
-    return prngs;
+
+    delete untwister;
+
+    return pyPRNGs;
 }
 
 /* Python interface */
@@ -128,7 +113,7 @@ BOOST_PYTHON_MODULE(untwister) {
     unsigned int threads = std::thread::hardware_concurrency();
     current.attr("THREADS") = threads;
 
-    def("prngs", Prngs, "\n List of supported PRNGs");
+    def("get_prngs", Prngs, "\n List of supported PRNGs");
 
     def(
         "find_seed",
@@ -138,11 +123,4 @@ BOOST_PYTHON_MODULE(untwister) {
         "\nThis is the cracking module for a generic mersenne twister 19932"
     );
 
-    def(
-        "generate_sample",
-        Sample,
-        (arg("prng"), arg("seed"), arg("depth") = 1000),
-        "\n Generate a sample using a given PRNG"
-    );
-
 }