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,3 +113,4 @@ venv.bak/
 *.model
 *.csv
 *.txt
+!requirements.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

@@ -0,0 +1,56 @@
+# Eyeballer
+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.
+
+## Installation
+
+To install in a python venv (Recommended):
+```
+python3 -m venv venv
+source venv/bin/activate
+make init
+```
+
+To install for Conda (automatically creates conda venv):
+
+```
+make conda
+```
+
+## Running Tests
+
+```
+make test
+```
+
+## Predicting Labels
+To eyeball some screenshots, just run the "predict" mode:
+
+```
+eyeballer --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
+```
+
+Or for a whole directory of files:
+
+```
+eyeballer --weights YOUR_WEIGHTS.h5 predict PATH_TO/YOUR_FILES/
+```
+
+## Training
+To train a new model, run:
+```
+eyeballer train
+```
+
+You'll want a machine with a good GPU for this to run in a reasonable amount of time. Setting that up is outside the scope of this readme, however.
+
+This will output a new model file (weights.h5 by default).
+
+## Evaluation
+
+You just trained a new model, cool! Let's see how well it performs against some images it's never seen before, across a variety of metrics:
+
+```
+eyeballer --weights YOUR_WEIGHTS.h5 evaluate
+```

eyeballer/__init__.py

@@ -1,2 +1 @@
-from eyeballer.eyeballer import EyeballModel  # noqa: F401
-from eyeballer.visualization import OcclusionGenerator  # noqa: F401
+from eyeballer.model import EyeballModel  # noqa: F401

eyeballer/eyeballer.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python3
+
 import click
 import csv
 
-from model import EyeballModel, DATA_LABELS
-from visualization import HeatMap
+from eyeballer.model import EyeballModel, DATA_LABELS
+from eyeballer.visualization import HeatMap
 
 
 @click.group(invoke_without_command=True)
@@ -64,8 +65,10 @@ def pretty_print_evaluation(results):
     for label in DATA_LABELS:
         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"All or nothing accuracy: {results['all_or_nothing_accuracy']:4.2%}")
     print(f"Top 10 worst predictions: {results['top_10_worst'][1]}")
 
 

eyeballer/model.py

@@ -1,7 +1,12 @@
 import os
 import random
+import sys
 
+# Prevent Tkinter Dependency
+import matplotlib
+matplotlib.use('agg')  # noqa: E402
 import matplotlib.pyplot as plt
+
 import numpy as np
 import pandas as pd
 import Augmentor
@@ -17,11 +22,24 @@ 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 augmentation import EyeballerAugmentation
+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.
 
@@ -33,7 +51,7 @@ class EyeballModel:
     image_width, image_height = 224, 224
     input_shape = (image_width, image_height, 3)
 
-    def __init__(self, print_summary=False, weights_file="weights.h5", seed=None):
+    def __init__(self, weights_file, print_summary=False, seed=None):
         """Constructor for model class.
 
         Keyword arguments:
@@ -78,10 +96,16 @@ class EyeballModel:
         self.training_labels = self.training_labels.sample(frac=1)
 
         if weights_file is not None and os.path.isfile(weights_file):
-            self.model.load_weights(weights_file)
+            try:
+                self.model.load_weights(weights_file)
+            except OSError:
+                print(f"ERROR: Unable to open weights file '{weights_file}'")
+                sys.exit(-1)
             print("Loaded model from file.")
         else:
-            print("No model loaded from file")
+            if weights_file is not None:
+                raise FileNotFoundError
+            print("WARN: No model loaded from file. Generating random model")
 
         # Data augmentation
         augmentor = Augmentor.Pipeline()
@@ -204,7 +228,7 @@ class EyeballModel:
             screenshots = os.listdir(path)
             screenshots = [os.path.join(path, s) for s in screenshots]
         else:
-            return None
+            raise FileNotFoundError
 
         results = []
         for screenshot in screenshots:
@@ -265,7 +289,22 @@ 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)
         predictions = predictions > threshold
         ground_truth = evaluation_generator.data
@@ -274,8 +313,38 @@ class EyeballModel:
         stats["all_or_nothing_accuracy"] = accuracy_score(ground_truth, predictions)
         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
 
+    def _none_of_the_above_recall(self, labels, predictions):
+        """Returns the recall score for the 'none of the above' images.
+        That means, all the images that don't have a category.
+        """
+        total_count = 0
+        correct_count = 0
+        for item in zip(labels.astype(int), predictions.astype(int)):
+            # Is this a none of the above?
+            if not item[0].any():
+                total_count += 1
+                if not item[1].any():
+                    correct_count += 1
+        return correct_count / total_count
+
+    def _none_of_the_above_precision(self, labels, predictions):
+        """Returns the precision score for the 'none of the above' images.
+        That means, all the images that don't have a category.
+        """
+        total_count = 0
+        correct_count = 0
+        for item in zip(labels.astype(int), predictions.astype(int)):
+            # Is this a none of the above prediction?
+            if not item[1].any():
+                total_count += 1
+                if not item[0].any():
+                    correct_count += 1
+        return correct_count / total_count
+
     def _top_images(self, labels, predictions, top_k=10, best=False):
         """Collect top-k best or top-k worst predicted images
 

eyeballer/visualization.py

@@ -3,7 +3,7 @@ import math
 from keras.preprocessing import image
 from copy import deepcopy
 from matplotlib import pyplot as plt
-from model import DATA_LABELS
+from eyeballer.model import DATA_LABELS
 
 
 class HeatMap():

setup.py

@@ -6,17 +6,28 @@ setup(
     packages=find_packages(exclude=["tests"]),
     url='https://github.com/BishopFox/eyeballer',
     license='GNU General Public License v3.0',
-    author='Daniel Petro',
+    author='Bishop Fox',
     author_email='',
     description='',
+    entry_points={
+        'console_scripts': [
+            'eyeballer = eyeballer.eyeballer:cli'
+        ]
+    },
     install_requires=[
-        'keras',
-        'sklearn',
+        'Augmentor',
+        'click',
+        'matplotlib',
         'numpy',
         'pandas',
-        'matplotlib',
-        'pillow', 'click'
+        'pillow',
+        'sklearn',
+        'keras @ git+https://github.com/keras-team/keras',
+    ],
+    extras_require={
+    'tf': ['tensorflow'],
+    'tf_gpu': ['tensorflow-gpu'],
+    },
+    dependency_links=[
     ],
-    test_requires=[],
-    python_requires='>=3.6',    # We use some python 3.6+ only features, like ordered dicts by default and f-strings
 )

tests/EvaluateTest.py

@@ -0,0 +1 @@
+#TODO

tests/PredictTest.py

@@ -0,0 +1,84 @@
+import errno
+import os
+import unittest
+
+from eyeballer.model import EyeballModel
+
+
+class PredictTest(unittest.TestCase):
+    def setUp(self):
+        weights_file = "latest_weights.h5"
+        if not os.path.isfile(weights_file):
+            print("Error: Symlink the latest weights file to 'latest_weights.h5'")
+            raise FileNotFoundError(
+                errno.ENOENT,
+                os.strerror(errno.ENOENT),
+                weights_file)
+
+        model_kwargs = {
+            "weights_file": weights_file,
+            "print_summary": False,
+            "seed": None
+        }
+        self.model = EyeballModel(**model_kwargs)
+
+    def test_different_seed_predict(self):
+        model_kwargs = {
+            "weights_file": None,
+            "print_summary": False,
+            "seed": 12345678
+        }
+        same_seed_model = EyeballModel(**model_kwargs)
+
+        screenshot = "tests/data/test.png"
+
+        results_one = self.model.predict(screenshot)
+        results_two = same_seed_model.predict(screenshot)
+
+        self.assertNotEqual(results_one, results_two)
+
+    def test_same_seed_predict(self):
+        model_kwargs = {
+            "weights_file": None,
+            "print_summary": False,
+            "seed": 12345678
+        }
+        same_seed_model = EyeballModel(**model_kwargs)
+
+        screenshot = "tests/data/test.png"
+
+        results_one = same_seed_model.predict(screenshot)
+        results_two = same_seed_model.predict(screenshot)
+
+        self.assertEquals(results_one, results_two)
+
+    def get_prediction_by_name(self, results):
+        prediction = None
+        max_val = 0.0
+
+        print(results)
+        for k, v in results[0].items():
+            if not isinstance(v, str) and v > max_val:
+                prediction = k
+                max_val = v
+        return prediction
+
+    def test_predict_custom404(self):
+        screenshot = "tests/data/custom404.png"
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "custom404")
+
+    def test_predict_login(self):
+        screenshot = "tests/data/login.png"
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "login")
+
+    def test_predict_homepage(self):
+        screenshot = "tests/data/homepage.png"
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "homepage")
+
+    def test_predict_oldlooking(self):
+        screenshot = "tests/data/oldlooking.png"
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "oldlooking")

tests/TrainTest.py

@@ -0,0 +1 @@
+#TODO

tests/__init__.py

@@ -0,0 +1 @@
+from tests.PredictTest import *