Re-org'd project, build-system, and added tests

- "--heatmap" flag under prediction

Makefile

@@ -0,0 +1,95 @@
+###############################################################################
+# 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)
+	$(PIP) install --upgrade --editable .
+
+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 +1 @@
-from eyeballer.eyeballer import EyeballModel  # noqa: F401
+from eyeballer.model import EyeballModel  # noqa: F401

eyeballer/eyeballer.py

@@ -1,8 +1,10 @@
 #!/usr/bin/env python3
+
 import click
 import csv
 
-from model import EyeballModel, DATA_LABELS
+from eyeballer.model import EyeballModel, DATA_LABELS
+from eyeballer.visualization import HeatMap
 
 
 @click.group(invoke_without_command=True)
@@ -34,10 +36,16 @@ 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.pass_context
-def predict(ctx, screenshot):
+def predict(ctx, screenshot, heatmap):
     model = EyeballModel(**ctx.obj['model_kwargs'])
     results = model.predict(screenshot)
+
+    if heatmap:
+        # Generate a heatmap
+        HeatMap(screenshot, model).generate()
+
     if not results:
         print("Error: Input file does not exist")
     if len(results) == 1:

eyeballer/model.py

@@ -3,7 +3,7 @@ import random
 
 # Prevent Tkinter Dependency
 import matplotlib
-matplotlib.use('agg')
+matplotlib.use('agg')  # noqa: E402
 import matplotlib.pyplot as plt
 
 import numpy as np
@@ -21,7 +21,7 @@ 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"]
@@ -189,6 +189,24 @@ class EyeballModel:
             plt.legend(['Train', 'Validation'], loc='upper left')
             plt.savefig(self.graphs_directory + "loss.png")
 
+    def predict_on_array(self, image):
+        """Predict the labels for a single screenshot
+
+        Keyword arguments:
+        image -- The numpy array of the image to classify
+        """
+        img = np.expand_dims(image, axis=0)
+        img = preprocess_input(img)
+
+        prediction = self.model.predict(img, batch_size=1)
+        result = dict()
+        result["filename"] = "custom-image"
+        result["custom404"] = prediction[0][0]
+        result["login"] = prediction[0][1]
+        result["homepage"] = prediction[0][2]
+        result["oldlooking"] = prediction[0][3]
+        return result
+
     def predict(self, path):
         """Predict the labels for a single file
 
@@ -203,7 +221,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:
@@ -267,7 +285,7 @@ class EyeballModel:
         prediction_manager = PredictionManager()
 
         try:
-            output = self.model.predict_generator(
+            self.model.predict_generator(
                 evaluation_generator,
                 callbacks=[prediction_manager],
                 verbose=1,

eyeballer/visualization.py

@@ -0,0 +1,127 @@
+import numpy as np
+import math
+from keras.preprocessing import image
+from copy import deepcopy
+from matplotlib import pyplot as plt
+from eyeballer.model import DATA_LABELS
+
+
+class HeatMap():
+    def __init__(self, screenshot_path, model, boxsize=28, step=7):
+        self.model = model
+
+        img = image.load_img(screenshot_path, target_size=(self.model.image_width, self.model.image_height))
+        img = image.img_to_array(img)
+
+        self.screenshot = img
+        self.boxsize = boxsize
+        self.step = step
+        self.x = 0
+        self.y = 0
+        self.screenshot_path = screenshot_path
+        self.gamma = 3
+
+    def generate(self, output_file="heatmap.png"):
+        """ Make a single heatmap image and return it """
+        heatmaps = []
+        labels = []
+        results = self.model.predict(self.screenshot_path)
+        results = results[0]
+        for label in DATA_LABELS:
+            boxsize = self.boxsize
+            # Ignore this label if it didn't predict positively (true label)
+            if results[label] > 0.5:
+                worst_score = 1
+                while worst_score > 0.5:
+                    heatmap, worst_score = self._get_heatmap(label, boxsize)
+                    if worst_score > 0.5:
+                        boxsize += 28
+                        print(f"Didn't get a good image for {label}. Trying again with a bigger boxsize: {boxsize}")
+                    else:
+                        heatmaps.append(heatmap)
+                        labels.append(label)
+
+        plt.figure()
+        screenshot_image = image.load_img(self.screenshot_path, target_size=(self.model.image_width, self.model.image_width))
+        _, subplots = plt.subplots(1, len(heatmaps))
+        # This is a little janky, but matplotlib returns a list above if there's multiple subplots, and just a single subplot if there's only one
+        #   it would have been easier if it was always a list. But alas
+        if not heatmaps:
+            print("No heatmap made. The image did not have a True classification for any label")
+            return
+        if len(heatmaps) > 1:
+            for i, heatmap in enumerate(heatmaps):
+                subplots[i].imshow(screenshot_image, cmap='binary', interpolation='none')
+                subplots[i].imshow(heatmap, alpha=0.5, interpolation='none')
+                subplots[i].set_title(labels[i])
+        else:
+            subplots.imshow(screenshot_image, cmap='binary', interpolation='none')
+            subplots.imshow(heatmap, alpha=0.5, interpolation='none')
+            subplots.set_title(labels[0])
+        plt.savefig(output_file)
+        print(f"Heatmap image written to: {output_file}")
+
+    def _get_heatmap(self, label, boxsize):
+        worst_score = 1
+        self.x, self.y = 0, 0
+
+        heatmap = np.ones((self.model.image_width, self.model.image_height))
+        heatmap *= 255
+        while True:
+            # Occlude an image
+            occluded_image, x, y = self._occlude(boxsize)
+            if occluded_image is not None:
+                # Score the new occluded image
+                results = self.model.predict_on_array(occluded_image)
+                score = results[label]  # TODO
+                worst_score = min(worst_score, score)
+                # Scale the score up to pixel values
+                score = math.floor(score * 255)
+                # Anneal the scores as you go outward into the occlusion block
+                new_scores = np.ones((boxsize, boxsize))
+                new_scores *= score
+                new_scores = self._gamma_anneal(new_scores)
+
+                # For each pixel, update only where the new score is lower
+                original_scores = heatmap[x: x+boxsize, y: y+boxsize]
+
+                target_area = np.where(new_scores < original_scores, new_scores, original_scores)
+                heatmap[x: x+boxsize, y: y+boxsize] = target_area
+            else:
+                break
+
+        return heatmap, worst_score
+
+    def _gamma_anneal(self, occlusion_area):
+        """ Get the annealed score based on centeredness of the occlusion area """
+        width, height = occlusion_area.shape
+
+        center_x = (width-1) / 2
+        center_y = (height-1) / 2
+
+        annealed_area = np.zeros(occlusion_area.shape)
+
+        for (x, y), score in np.ndenumerate(occlusion_area):
+            distance = math.sqrt((center_x - x)**2 + (center_y - y)**2)
+            annealed_area[x, y] = min(255, score + (self.gamma * distance))
+
+        return annealed_area
+
+    def _occlude(self, boxsize):
+        ''' Return a single occluded image and its location '''
+        if self.x + boxsize > self.screenshot.shape[0]:
+            return None, None, None
+
+        retImg = np.copy(self.screenshot)
+        retImg[self.x:self.x+boxsize, self.y:self.y+boxsize] = 0.0
+
+        old_i = deepcopy(self.x)
+        old_j = deepcopy(self.y)
+
+        # update indices
+        self.y = self.y + self.step
+        if self.y+boxsize > self.screenshot.shape[1]:  # reached end
+            self.y = 0  # reset j
+            self.x = self.x + self.step  # go to next row
+
+        return retImg, old_i, old_j

requirements.txt

@@ -1,9 +0,0 @@
-Augmentor
-click
-matplotlib
-numpy
-pandas
-pillow
-sklearn
-tensorflow
-git+https://github.com/keras-team/keras/

setup.py

@@ -0,0 +1,30 @@
+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='Daniel Petro',
+    author_email='',
+    description='',
+    entry_points={
+        'console_scripts': [
+            'eyeballer = eyeballer.eyeballer:cli'
+        ]
+    },
+    install_requires=[
+        'Augmentor',
+        'click',
+        'matplotlib',
+        'numpy',
+        'pandas',
+        'pillow',
+        'sklearn',
+        'tensorflow',
+        'keras @ git+https://github.com/keras-team/keras',
+    ],
+    dependency_links=[
+    ],
+)

tests/EvaluateTest.py

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

tests/PredictTest.py

@@ -0,0 +1,83 @@
+import errno
+import os
+import unittest
+
+from eyeballer.model import EyeballModel, DATA_LABELS
+
+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 *