Added a GPU check for NVIDIA in Makefile

This fix will only work on Linux. We would need to use 'ioreg' for Mac

.gitignore

@@ -113,6 +113,4 @@ venv.bak/
 *.model
 *.csv
 *.txt
-results.html
 !requirements.txt
-!requirements-gpu.txt

Makefile

@@ -0,0 +1,99 @@
+###############################################################################
+# Parts of this are borrowed from the HTTPie Makefile. Thanks!
+###############################################################################
+
+REQUIREMENTS="requirements-dev.txt"
+TAG="\n\n\033[0;32m\#\#\# "
+END=" \#\#\# \033[0m\n"
+
+PIP=$(shell which pip3)
+CONDA=$(shell which conda)
+CONDA_HOME="$(HOME)/anaconda3"
+VENV="eyeballer_venv"
+
+SHELL := /bin/bash
+.ONE_SHELL:
+
+all: test
+
+# Target-specific assignments
+conda: PIP := "$(CONDA_HOME)/envs/$(VENV)/bin/pip"
+conda init: uninstall-eyeballer
+
+ifeq ($(strip $(MAKECMDGOALS)),conda)
+	@echo -e $(TAG)Conda detected at: $(CONDA)$(END)
+	conda create -yn $(VENV) | true
+
+	source $(CONDA_HOME)/etc/profile.d/conda.sh; \
+	conda activate $(VENV); \
+	conda install -y pip
+endif
+	@echo -e $(TAG)Installing dev requirements$(END)
+	$(PIP) install --upgrade -r $(REQUIREMENTS)
+
+	@echo -e $(TAG)Installing Eyeballer$(END)
+ifeq ($(shell lspci | grep NVIDIA),NVIDIA))
+	$(PIP) install --upgrade --editable .[tf_gpu]
+else
+	$(PIP) install --upgrade --editable .[tf]
+endif
+
+ifeq ($(strip $(MAKECMDGOALS)),conda)
+	@echo -e $(TAG)Activate virtualenv to run eyeballer: $(END)
+	@echo conda activate $(VENV)
+endif
+
+	@echo
+
+clean:
+	@echo -e $(TAG)Cleaning up$(END)
+	rm -rf *.egg dist build .cache .pytest_cache eyeballer.egg-info
+	find . -name '__pycache__' -delete -print -o -name '*.pyc' -delete -print
+ifneq ($(strip $(CONDA)),)
+	conda env remove -y --name eyeballer_venv | true
+endif
+	@echo
+
+
+###############################################################################
+# Testing
+###############################################################################
+
+
+test:
+	@echo -e $(TAG)Running tests on the current Python interpreter with coverage $(END)
+	python -m unittest tests
+	@echo
+
+
+# test-all is meant to test everything — even this Makefile
+test-all: uninstall-all clean init test pycodestyle
+	@echo
+
+
+pycodestyle:
+	which pycodestyle || $(PIP) install pycodestyle
+	pycodestyle
+	@echo
+
+
+###############################################################################
+# Uninstalling
+###############################################################################
+
+uninstall-eyeballer:
+	@echo using pip3 located at: $(PIP)
+	@echo -e $(TAG)Uninstalling eyeballer$(END)
+	- $(PIP) uninstall --yes eyeballer &2>/dev/null
+
+	@echo "Verifying…"
+	cd .. && ! python -m eyeballer --version &2>/dev/null
+
+	@echo "Done"
+	@echo
+
+
+uninstall-all: uninstall-eyeballer
+
+	@echo -e $(TAG)Uninstalling development requirements$(END)
+	- $(PIP) uninstall --yes -r $(REQUIREMENTS)

README.md

@@ -3,16 +3,25 @@ Give those screenshots of yours a quick eyeballing.
 
 Eyeballer is meant for large-scope network penetration tests where you need to find "interesting" targets from a huge set of web-based hosts. Go ahead and use your favorite screenshotting tool like normal (EyeWitness or GoWitness) and then run them through Eyeballer to tell you what's likely to contain vulnerabilities, and what isn't.
 
-## Setup
+## Installation
 
-Download required packages on pip:
+To install in a python venv (Recommended):
 ```
-sudo pip3 install -r requirements.txt
+python3 -m venv venv
+source venv/bin/activate
+make init
 ```
 
-Or if you want GPU support:
+To install for Conda (automatically creates conda venv):
+
 ```
-sudo pip3 install -r requirements-gpu.txt
+make conda
+```
+
+## Running Tests
+
+```
+make test
 ```
 
 ## Predicting Labels

eyeballer/__init__.py

@@ -0,0 +1 @@
+from eyeballer.model import EyeballModel  # noqa: F401

eyeballer/eyeballer.py

@@ -5,7 +5,6 @@ import csv
 
 from eyeballer.model import EyeballModel, DATA_LABELS
 from eyeballer.visualization import HeatMap
-from jinja2 import Template
 
 
 @click.group(invoke_without_command=True)
@@ -38,9 +37,8 @@ def train(ctx, graphs, batchsize, epochs):
 @cli.command()
 @click.argument('screenshot')
 @click.option('--heatmap', default=False, is_flag=True, help="Create a heatmap graphfor the prediction")
-@click.option('--threshold', default=.5, type=float, help="Threshold confidence for labeling")
 @click.pass_context
-def predict(ctx, screenshot, heatmap, threshold):
+def predict(ctx, screenshot, heatmap):
     model = EyeballModel(**ctx.obj['model_kwargs'])
     results = model.predict(screenshot)
 
@@ -58,61 +56,20 @@ def predict(ctx, screenshot, heatmap, threshold):
             labelwriter = csv.DictWriter(csvfile, fieldnames=fieldnames)
             labelwriter.writeheader()
             labelwriter.writerows(results)
-
         print("Output written to results.csv")
-        buildHTML(processResults(results))
-        print("HTML written to results.html")
-
-
-def processResults(results):
-    '''Filter the initial results dictionary and reformat it for use in JS.
-
-        Keyword arguments:
-        results -- dictionary output from predict function
-
-    '''
-    jsResults = {}
-
-    for result in results:
-        positiveTags = []
-
-        for label, label_info in result.items():
-            if (label == 'filename'):
-                pass
-
-            elif label_info:
-                positiveTags.append(label)
-
-        jsResults[result['filename']] = positiveTags
-    return(jsResults)
-
-
-def buildHTML(jsResults):
-    '''Build HTML around the JS Dictionary that is passed from processResults.
-
-        Keyword arguments:
-        jsResults -- dictionary output from processResults function
-    '''
-    html_output = ""
-    with open("prediction_output_template.html") as template_file:
-        template = Template(template_file.read())
-        html_output = template.render(jsResults=jsResults)
-
-    with open('results.html', 'w') as file:
-        file.write(html_output)
 
 
 def pretty_print_evaluation(results):
     """Print a human-readable summary of the evaluation"""
     # We use 4.2% to handle all the way from "  0.00%" (7chars) to "100.00%" (7chars)
     for label in DATA_LABELS:
-        print("{} Precision Score: {:4.2%}".format(label, results[label]['precision']))
-        print("{} Recall Score: {:4.2%}".format(label, results[label]['recall']))
-    print("'None of the above' Precision: {:4.2%}".format(results['none_of_the_above_precision']))
-    print("'None of the above' Recall: {:4.2%}".format(results['none_of_the_above_recall']))
-    print("All or nothing Accuracy: {:4.2%}".format(results['all_or_nothing_accuracy']))
-    print("Overall Binary Accuracy: {:4.2%}".format(results['total_binary_accuracy']))
-    print("Top 10 worst predictions: {}".format(results['top_10_worst'][1]))
+        print(f"{label} Precision Score: {results[label]['precision']:4.2%}")
+        print(f"{label} Recall Score: {results[label]['recall']:4.2%}")
+    print(f"'None of the above' Precision: {results['none_of_the_above_precision']:4.2%}")
+    print(f"'None of the above' Recall: {results['none_of_the_above_recall']:4.2%}")
+    print(f"All or nothing Accuracy: {results['all_or_nothing_accuracy']:4.2%}")
+    print(f"Overall Binary Accuracy: {results['total_binary_accuracy']:4.2%}")
+    print(f"Top 10 worst predictions: {results['top_10_worst'][1]}")
 
 
 @cli.command()

eyeballer/model.py

@@ -19,13 +19,27 @@ from keras.layers import Dropout, Dense
 from keras.layers import GlobalAveragePooling2D
 from keras.models import Model
 from keras.optimizers import Adam
+from keras.preprocessing import image
 from keras.preprocessing.image import ImageDataGenerator
 from sklearn.metrics import classification_report, accuracy_score, hamming_loss
 from eyeballer.augmentation import EyeballerAugmentation
+from keras.callbacks import Callback
 
 DATA_LABELS = ["custom404", "login", "homepage", "oldlooking"]
 
 
+class PredictionManager(Callback):
+    def __init__(self):
+        super(PredictionManager, self).__init__()
+        self.predictions = []
+
+    def getPredictions(self):
+        return np.array(self.predictions)
+
+    def on_predict_batch_end(self, batch, logs=None):
+        self.predictions.append(logs['outputs'][0].tolist()[0])
+
+
 class EyeballModel:
     """The primary model class of Eyeballer.
 
@@ -77,7 +91,7 @@ class EyeballModel:
             self.random_seed = True
             self.seed = random.randint(0, 999999)
             print("No seed set, ", end='')
-        print("using seed: {}".format(self.seed))
+        print(f"using seed: {self.seed}")
         random.seed(self.seed)
         self.training_labels = self.training_labels.sample(frac=1)
 
@@ -85,7 +99,7 @@ class EyeballModel:
             try:
                 self.model.load_weights(weights_file)
             except OSError:
-                print("ERROR: Unable to open weights file '{}'".foramt(weights_file))
+                print(f"ERROR: Unable to open weights file '{weights_file}'")
                 sys.exit(-1)
             print("Loaded model from file.")
         else:
@@ -200,12 +214,11 @@ class EyeballModel:
         result["oldlooking"] = prediction[0][3]
         return result
 
-    def predict(self, path, threshold=0.5):
+    def predict(self, path):
         """Predict the labels for a single file
 
         Keyword arguments:
-        path -- The path to the file that we'll be evaluating.
-            NOTE: This actually has to be a directory containing a directory with our files. Because Keras is dumb. Don't ask.
+        path -- The path to the file that we'll be evaluating
         """
         # Is this a single file, or a directory?
         screenshots = []
@@ -217,30 +230,23 @@ class EyeballModel:
         else:
             raise FileNotFoundError
 
-        # Prepare the data
-        data_generator = ImageDataGenerator(
-            preprocessing_function=preprocess_input)
-        prediction_generator = data_generator.flow_from_directory(
-            directory=path,
-            target_size=(self.image_width, self.image_height),
-            shuffle=False,
-            class_mode=None)
-
-        predictions = self.model.predict_generator(
-            prediction_generator,
-            verbose=1,
-            steps=len(prediction_generator))
-
-        predictions = predictions > threshold
         results = []
-        for filename, prediction in zip(prediction_generator.filenames, predictions):
+        for screenshot in screenshots:
+            # Load the image into memory
+            img = image.load_img(screenshot, target_size=(self.image_width, self.image_height))
+            img = image.img_to_array(img)
+            img = np.expand_dims(img, axis=0)
+            img = preprocess_input(img)
+
+            prediction = self.model.predict(img, batch_size=1)
             result = dict()
-            result["filename"] = path + filename
-            result["custom404"] = prediction[0]
-            result["login"] = prediction[1]
-            result["homepage"] = prediction[2]
-            result["oldlooking"] = prediction[3]
+            result["filename"] = screenshot
+            result["custom404"] = prediction[0][0]
+            result["login"] = prediction[0][1]
+            result["homepage"] = prediction[0][2]
+            result["oldlooking"] = prediction[0][3]
             results.append(result)
+
         return results
 
     def evaluate(self, threshold=0.5):
@@ -283,21 +289,30 @@ class EyeballModel:
         else:
             print("Using evaluation set...")
 
-        predictions = self.model.predict_generator(
-            evaluation_generator,
-            verbose=1,
-            steps=len(evaluation_generator))
+        prediction_manager = PredictionManager()
+
+        try:
+            self.model.predict_generator(
+                evaluation_generator,
+                callbacks=[prediction_manager],
+                verbose=1,
+                steps=len(evaluation_generator))
+        except KeyboardInterrupt:
+            print('\r\nJob interrupted (Ctrl+C): Outputting completed results')
+
+        predictions = prediction_manager.getPredictions()
+
+        # Truncates evaluation generator if test was interrupted
+        evaluation_generator._data = evaluation_generator.data[:predictions.shape[0]]
 
         self._save_prediction_histograms(predictions)
-        print("predictions: ", type(predictions))
         predictions = predictions > threshold
-        ground_truth = self.evaluation_labels[DATA_LABELS].to_numpy()
-        filenames = self.evaluation_labels[["filename"]].to_numpy()
+        ground_truth = evaluation_generator.data
         stats = classification_report(ground_truth, predictions, target_names=DATA_LABELS, output_dict=True)
         stats["total_binary_accuracy"] = 1 - hamming_loss(ground_truth, predictions)
         stats["all_or_nothing_accuracy"] = accuracy_score(ground_truth, predictions)
-        stats["top_10_best"] = self._top_images(filenames, ground_truth, predictions, best=True)
-        stats["top_10_worst"] = self._top_images(filenames, ground_truth, predictions, best=False)
+        stats["top_10_best"] = self._top_images(evaluation_generator, predictions, best=True)
+        stats["top_10_worst"] = self._top_images(evaluation_generator, predictions, best=False)
         stats["none_of_the_above_recall"] = self._none_of_the_above_recall(ground_truth, predictions)
         stats["none_of_the_above_precision"] = self._none_of_the_above_precision(ground_truth, predictions)
         return stats
@@ -314,9 +329,6 @@ class EyeballModel:
                 total_count += 1
                 if not item[1].any():
                     correct_count += 1
-        if total_count == 0:
-            print("WARNING: None of the Above Recall is NaN")
-            return 0
         return correct_count / total_count
 
     def _none_of_the_above_precision(self, labels, predictions):
@@ -331,23 +343,20 @@ class EyeballModel:
                 total_count += 1
                 if not item[0].any():
                     correct_count += 1
-        if total_count == 0:
-            print("WARNING: None of the Above Precision is NaN")
-            return 0
         return correct_count / total_count
 
-    def _top_images(self, filenames, ground_truth, predictions, top_k=10, best=False):
+    def _top_images(self, labels, predictions, top_k=10, best=False):
         """Collect top-k best or top-k worst predicted images
 
         Keyword arguments:
-        ground_truth -- The correct labels
+        labels -- The keras generator that contains the filenames and data labels
         predictions -- The numpy array of predictions
         top_k -- Top k elements
         best -- True/False. Calculate either the best or worst images
 
         :Return -- Tuple of top-k indicies and top-k filenames
         """
-        true_labels = np.array(ground_truth).astype(float)
+        true_labels = np.array(labels.data).astype(float)
         predictions = predictions.astype(float)
 
         differences = np.abs(predictions - true_labels).sum(axis=1)
@@ -360,7 +369,7 @@ class EyeballModel:
             indicies = np.flipud(indicies)
 
         for i in indicies[:top_k]:
-            top_file_list.append(filenames[i][0])
+            top_file_list.append(labels.filenames[i])
 
         return indicies[:top_k], top_file_list
 

eyeballer/visualization.py

@@ -36,7 +36,7 @@ class HeatMap():
                     heatmap, worst_score = self._get_heatmap(label, boxsize)
                     if worst_score > 0.5:
                         boxsize += 28
-                        print("Didn't get a good image for {}. Trying again with a bigger boxsize: {}".format(label, boxsize))
+                        print(f"Didn't get a good image for {label}. Trying again with a bigger boxsize: {boxsize}")
                     else:
                         heatmaps.append(heatmap)
                         labels.append(label)
@@ -59,7 +59,7 @@ class HeatMap():
             subplots.imshow(heatmap, alpha=0.5, interpolation='none')
             subplots.set_title(labels[0])
         plt.savefig(output_file)
-        print("Heatmap image written to: {}".format(output_file))
+        print(f"Heatmap image written to: {output_file}")
 
     def _get_heatmap(self, label, boxsize):
         worst_score = 1

prediction_output_template.html

@@ -1,189 +0,0 @@
-<!DOCTYPE HTML>
-<meta charset="utf-8"/>
-<html>
-<title>Eyeballer Results</title>
-<body>
-    <div id="btnsCont">
-      <button class="btn all" onclick="filterSelection('all')"> Show all</button>
-      <button class="btn" onclick="filterSelection('custom404')"> 404</button>
-      <button class="btn" onclick="filterSelection('login')"> Login</button>
-      <button class="btn" onclick="filterSelection('homepage')"> Home Page</button>
-      <button class="btn" onclick="filterSelection('oldlooking')"> Old Looking</button>
-        <select onchange="numOnPage=this.value; filterSelection([]);" style="float:right; margin:5px; margin-right:30px; height:30px; width:auto;">
-          <option value="25">25</option>
-          <option value="50">50</option>
-          <option value="100">100</option>
-          <option value="250">250</option>
-        </select>
-    </div>
-    <div id="btnsPage">
-    </div>
-    <div id="container">
-    </div>
-    <button onclick="topFunction()" id="topBtn" title="Go to top">Top</button>
-<script>
-    var results = {{jsResults}};
-    var selectedTags = new Array();
-    var btnContainer = document.getElementById("btnsCont");
-    var categoryBtns = btnContainer.getElementsByClassName("btn");
-    var numOnPage = 100;
-    var indexesToDisplay = [];
-    filterSelection([]);
-    for (var i = 0; i < categoryBtns.length; i++) {
-        categoryBtns[i].addEventListener("click", function () {
-            if (this.className.indexOf("active") > -1) {
-                this.className = this.className.replace(" active", "");
-            }
-            else {
-            this.className += " active";
-            }
-        });
-    }
-    function filterSelection(c) {
-        indexesToDisplay = [];
-        var currentIndex = -1;
-        if (c == "all") location.reload();
-        if (selectedTags.includes(c)!=true) selectedTags=selectedTags.concat(c);
-        else selectedTags.splice(selectedTags.indexOf(c),1);
-        for (var filename in results){
-                currentIndex++;
-                var tags = results[filename];
-                var isMissingTag = false;
-                for (var i in selectedTags){
-                    if (!tags.includes(selectedTags[i])){
-                        isMissingTag = true;
-                        break;
-                    }
-                }
-                if (isMissingTag === false){
-                    indexesToDisplay = indexesToDisplay.concat(currentIndex);
-                }
-            }
-        buildPageButtons(indexesToDisplay);
-        buildHTML(0, numOnPage);
-    }
-    function buildPageButtons(indexes){
-        numOfPages = Math.ceil(indexes.length/numOnPage);
-        var html = '';
-        for (var i = 0; i < numOfPages; i++){
-            html += ('<button class="btnpage" onclick="buildHTML(' + (i*numOnPage) + ', ' + ((i+1)*numOnPage) + '); activatePage(this);">' + (i+1) +'</button>');
-        }
-        document.getElementById('btnsPage').innerHTML = html;
-    }
-    function activatePage(pressedBtn){
-        var activeNow = document.getElementsByClassName('activePage')
-        for (var i = 0; i < activeNow.length; i++){
-            activeNow[i].className = activeNow[i].className.replace(" activePage", "");
-        }
-        pressedBtn.className += ' activePage';
-    }
-    function buildHTML(rangeMin, rangeMax){
-        var html = '';
-        if (rangeMax > indexesToDisplay.length) rangeMax=indexesToDisplay.length;
-        for (var i = rangeMin; i < rangeMax; i++){
-            index=indexesToDisplay[i];
-            filename = Object.keys(results)[index];
-            html += ('<div class="filterDiv"><a href=' + filename + ' target="_blank"> ' + results[filename].toString().replace(/,/g, ' ') + '</a><img class="res_pic" src="' + filename + '"></div>');
-        }
-        document.getElementById('container').innerHTML = html;
-    }
-    window.onscroll = function() {scrollFunction()};
-    function scrollFunction() {
-      if (document.body.scrollTop > 20 || document.documentElement.scrollTop > 20) {
-        document.getElementById("topBtn").style.display = "block";
-      } else {
-        document.getElementById("topBtn").style.display = "none";
-      }
-    }
-    function topFunction() {
-      document.body.scrollTop = 0; // For Safari
-      document.documentElement.scrollTop = 0; // For Chrome, Firefox, IE and Opera
-    }
-</script>
-    <style>
-    #topBtn {
-      display: none; /* Hidden by default */
-      position: fixed; /* Fixed/sticky position */
-      bottom: 20px; /* Place the button at the bottom of the page */
-      right: 30px; /* Place the button 30px from the right */
-      z-index: 99; /* Make sure it does not overlap */
-      border: none; /* Remove borders */
-      outline: none; /* Remove outline */
-      background-color: #193b72; /* Set a background color */
-      color: white; /* Text color */
-      cursor: pointer; /* Add a mouse pointer on hover */
-      padding: 15px; /* Some padding */
-      border-radius: 10px; /* Rounded corners */
-      font-size: 18px; /* Increase font size */
-      font-weight:bold;
-      border:solid;
-    }
-    #topBtn:hover {
-      background-color: #555;
-    }
-    .container {
-    overflow: hidden;
-    }
-    #btnsPage {
-    margin:10px;
-    }
-    .btnpage{
-        width:40px;
-        height:30px;
-        background-color: #f1f1f1;
-        outline: none;
-    }
-    .res_pic{
-    width:100%;
-    height:auto;
-    }
-    a{
-    font-weight:600;
-    padding:3px;
-    text-decoration: none;
-    color:#ffffff;
-    font-size:16px;
-    font-family:Sans-serif;
-    }
-    .name{
-    float:left;
-    }
-    .filterDiv {
-      float: left;
-      background-color: #777;
-      color: #ffffff;
-      width: 22%;
-      border-style:solid;
-      border-color:#666;
-      text-align: center;
-      margin: 10px;
-      line-height: 30px;
-      display: block;
-    }
-    .btn {
-      border: none;
-      outline: none;
-      padding: 12px 16px;
-      margin-left:10px;
-      background-color: #f1f1f1;
-      cursor: pointer;
-    }
-    .btn:hover {
-      background-color: #ddd;
-    }
-    .active {
-      background-color: #666;
-      color: white;
-    }
-    .activePage{
-        border: solid;
-        border-color:#666666;
-        color: black;
-    }
-    .btn.all {
-      background-color: #193b72;
-      color: white;
-    }
-    </style>
-</body>
-</html>

requirements-gpu.txt

@@ -1,9 +0,0 @@
-Augmentor
-click
-matplotlib
-numpy
-pandas
-pillow
-sklearn
-tensorflow-gpu
-keras

requirements.txt

@@ -1,10 +0,0 @@
-Augmentor
-click
-matplotlib
-numpy
-pandas
-pillow
-sklearn
-tensorflow
-keras
-jinja2

setup.py

@@ -0,0 +1,33 @@
+from setuptools import find_packages, setup
+
+setup(
+    name='eyeballer',
+    version='0.0.1',
+    packages=find_packages(exclude=["tests"]),
+    url='https://github.com/BishopFox/eyeballer',
+    license='GNU General Public License v3.0',
+    author='Bishop Fox',
+    author_email='',
+    description='',
+    entry_points={
+        'console_scripts': [
+            'eyeballer = eyeballer.eyeballer:cli'
+        ]
+    },
+    install_requires=[
+        'Augmentor',
+        'click',
+        'matplotlib',
+        'numpy',
+        'pandas',
+        'pillow',
+        'sklearn',
+        'keras @ git+https://github.com/keras-team/keras',
+    ],
+    extras_require={
+    'tf': ['tensorflow'],
+    'tf_gpu': ['tensorflow-gpu'],
+    },
+    dependency_links=[
+    ],
+)