Updated requirements

The release of TF2.0 and Numpy 1.17 causes a number of "future warning" error messages. Specifically setting things to include Keras and Numpy 1.16.4 removes these errors without needing to update the entire code base.

.gitignore

@@ -116,3 +116,6 @@ venv.bak/
 results.html
 !requirements.txt
 !requirements-gpu.txt
+!tests/models/*
+!tests/data/*
+!docs/*

README.md

@@ -3,8 +3,21 @@ 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.
 
+
+### Example Labels
+
+| Old-Looking Sites | Login Pages |
+| ------ |:-----:|
+| ![Sample Old-looking Page](/docs/old-looking.png) | ![Sample Login Page](/docs/login.png) |
+
+| Homepages | Custom 404's |
+| ------ |:-----:|
+| ![Sample Homepage](/docs/homepage.png) | ![Sample Custom 404](/docs/404.png) |
+
 ## Setup
 
+Eyeballer uses TF.keras on Tensorflow 2.0. This is (as of this moment) still in "beta". So the pip requirement for it looks a bit weird. It'll also probably conflict with an existing TensorFlow installation if you've got the regular 1.0 version installed. So, heads-up there. But 2.0 should be out of beta and official "soon" according to Google, so this problem ought to solve itself in short order.
+
 Download required packages on pip:
 ```
 sudo pip3 install -r requirements.txt
@@ -15,23 +28,40 @@ Or if you want GPU support:
 sudo pip3 install -r requirements-gpu.txt
 ```
 
+**NOTE**: Setting up a GPU for use with TensorFlow is way beyond the scope of this README. There's hardware compatibility to consider, drivers to install... There's a lot. So you're just going to have to figure this part out on your own if you want a GPU. But at least from a Python package perspective, the above requirements file has you covered.
+
+**Training Data** You can find our training data here:
+
+https://www.dropbox.com/sh/7aouywaid7xptpq/AAD_-I4hAHrDeiosDAQksnBma?dl=1
+
+Pretty soon, we're going to add this as a TensorFlow DataSet, so you don't need to download this separately like this. It'll also let us version the data a bit better. But for now, just deal with it. There's two things you need from the training data:
+
+1. `images/` folder, containing all the screenshots (resized down to 224x140. We'll have the full-size images up soon)
+2. `labels.csv` that has all the labels
+3. `bishop-fox-pretrained-v1.h5` A pretrained weights file you can use right out of the box without training.
+
+Copy all three into the root of the Eyeballer code tree.
+
 ## Predicting Labels
 To eyeball some screenshots, just run the "predict" mode:
 
 ```
-eyeballer --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
+eyeballer.py --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/
+eyeballer.py --weights YOUR_WEIGHTS.h5 predict PATH_TO/YOUR_FILES/
 ```
 
+Eyeballer will spit the results back to you in human readable format (a `results.html` file so you can browse it easily) and machine readable format (a `results.csv` file).
+
+
 ## Training
 To train a new model, run:
 ```
-eyeballer train
+eyeballer.py 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.
@@ -43,5 +73,7 @@ This will output a new model file (weights.h5 by default).
 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.py --weights YOUR_WEIGHTS.h5 evaluate
 ```
+
+The output will describe the model's accuracy in both recall and precision for each of the program's labels. (Including "none of the above" as a pseudo-label)

eyeballer.py

@@ -46,7 +46,7 @@ def predict(ctx, screenshot, heatmap, threshold):
 
     if heatmap:
         # Generate a heatmap
-        HeatMap(screenshot, model).generate()
+        HeatMap(screenshot, model, threshold).generate()
 
     if not results:
         print("Error: Input file does not exist")
@@ -60,11 +60,11 @@ def predict(ctx, screenshot, heatmap, threshold):
             labelwriter.writerows(results)
 
         print("Output written to results.csv")
-        buildHTML(processResults(results))
+        buildHTML(processResults(results, threshold))
         print("HTML written to results.html")
 
 
-def processResults(results):
+def processResults(results, threshold):
     '''Filter the initial results dictionary and reformat it for use in JS.
 
         Keyword arguments:
@@ -79,8 +79,7 @@ def processResults(results):
         for label, label_info in result.items():
             if (label == 'filename'):
                 pass
-
-            elif label_info:
+            elif label_info > threshold:
                 positiveTags.append(label)
 
         jsResults[result['filename']] = positiveTags

eyeballer/augmentation.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from Augmentor.Operations import Operation
-from keras.applications.mobilenet import preprocess_input
+from tensorflow.keras.applications.mobilenet import preprocess_input
 
 
 class EyeballerAugmentation(Operation):

eyeballer/model.py

@@ -1,6 +1,7 @@
 import os
 import random
 import sys
+import progressbar
 
 # Prevent Tkinter Dependency
 import matplotlib
@@ -10,16 +11,9 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import Augmentor
+import tensorflow as tf
 
-from keras import regularizers
-from keras.applications import MobileNet
-from keras.applications.mobilenet import preprocess_input
-from keras.callbacks import ModelCheckpoint
-from keras.layers import Dropout, Dense
-from keras.layers import GlobalAveragePooling2D
-from keras.models import Model
-from keras.optimizers import Adam
-from keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.applications.mobilenet import preprocess_input
 from sklearn.metrics import classification_report, accuracy_score, hamming_loss
 from eyeballer.augmentation import EyeballerAugmentation
 
@@ -33,11 +27,11 @@ class EyeballModel:
     """
     graphs_directory = "graphs/"
     checkpoint_file = "weights.h5"
-    image_dir = "images_224x224/"
+    image_dir = "images/"
     image_width, image_height = 224, 224
     input_shape = (image_width, image_height, 3)
 
-    def __init__(self, weights_file, print_summary=False, seed=None):
+    def __init__(self, weights_file, print_summary=False, seed=None, quiet=False):
         """Constructor for model class.
 
         Keyword arguments:
@@ -45,26 +39,29 @@ class EyeballModel:
         weights_file -- A filename for where to load the model's weights.
         seed -- PRNG seed, useful for repeating a previous run and using the same data. Training/Validation split is determined randomly.
         """
-        # Build the model
-        base_model = MobileNet(weights='imagenet', include_top=False, input_shape=self.input_shape)
-        x = base_model.output
-        x = GlobalAveragePooling2D()(x)
-        x = Dense(256, activation="relu", name="HiddenLayer1")(x)
-        x = Dropout(0.5)(x)
-        x = Dense(128, activation="relu", name="HiddenLayer2", kernel_regularizer=regularizers.l2(0.01))(x)
-        x = Dropout(0.2)(x)
-        output_layer = Dense(len(DATA_LABELS), activation="sigmoid", name="OutputLayer")(x)
-        self.model = Model(inputs=base_model.input, outputs=output_layer)
+        # # Build the model
+        self.model = tf.keras.Sequential()
+        pretrained_layer = tf.keras.applications.mobilenet.MobileNet(weights='imagenet', include_top=False, input_shape=self.input_shape)
+        self.model.add(pretrained_layer)
+        self.model.add(tf.keras.layers.GlobalAveragePooling2D())
+        self.model.add(tf.keras.layers.Dense(256, activation="relu"))
+        self.model.add(tf.keras.layers.Dropout(0.5))
+        self.model.add(tf.keras.layers.Dense(128, activation="relu"))
+        self.model.add(tf.keras.layers.Dropout(0.2))
+        self.model.add(tf.keras.layers.Dense(len(DATA_LABELS), activation="sigmoid"))
 
-        # for layer in base_model.layers:
+        self.model.compile(optimizer=tf.keras.optimizers.Adam(0.0005),
+                           loss="binary_crossentropy",
+                           metrics=["accuracy"])
+
+        # for layer in pretrained_layer.layers:
         #     layer.trainable = False
 
-        adam = Adam(lr=0.0005)
-        self.model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
-
         if print_summary:
             print(self.model.summary())
 
+        self.quiet = quiet
+
         # Pull out our labels for use in generators later
         data = pd.read_csv("labels.csv")
         self.training_labels = data.loc[data['evaluation'] == False]  # noqa: E712
@@ -117,7 +114,7 @@ class EyeballModel:
         """
         print("Training with seed: " + str(self.seed))
 
-        data_generator = ImageDataGenerator(
+        data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
             preprocessing_function=self.preprocess_training_function,
             validation_split=0.2,
             samplewise_center=True)
@@ -146,8 +143,12 @@ class EyeballModel:
             class_mode="other")
 
         # Model checkpoint - Saves model weights when validation accuracy improves
-        callbacks = [ModelCheckpoint(self.checkpoint_file, monitor='val_loss', verbose=1, save_best_only=True,
-                                     save_weights_only=True, mode='min')]
+        callbacks = [tf.keras.callbacks.ModelCheckpoint(self.checkpoint_file,
+                     monitor='val_loss',
+                     verbose=1,
+                     save_best_only=True,
+                     save_weights_only=True,
+                     mode='min')]
 
         history = self.model.fit_generator(
             training_generator,
@@ -201,11 +202,10 @@ class EyeballModel:
         return result
 
     def predict(self, path, threshold=0.5):
-        """Predict the labels for a single file
+        """Predict the labels for a single file or directory of files
 
         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(s) that we'll be evaluating.
         """
         # Is this a single file, or a directory?
         screenshots = []
@@ -217,29 +217,32 @@ 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):
+        bar = progressbar.ProgressBar()
+        if self.quiet:
+            bar = progressbar.NullBar()
+        for screenshot in bar(screenshots):
+            # Load the image into memory
+            img = None
+            try:
+                img = tf.keras.preprocessing.image.load_img(screenshot, target_size=(self.image_width, self.image_height))
+                img = tf.keras.preprocessing.image.img_to_array(img)
+                img = np.expand_dims(img, axis=0)
+                img = preprocess_input(img)
+            except IsADirectoryError:
+                print("\nWARN: Skipping directory: ", screenshot)
+                continue
+            except OSError:
+                print("\nWARN: Skipping empty or corrupt file: ", screenshot)
+                continue
+
+            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
 
@@ -250,7 +253,7 @@ class EyeballModel:
         threshold -- Value between 0->1. The cutoff where the numerical prediction becomes boolean. Default: 0.5
         """
         # Prepare the data
-        data_generator = ImageDataGenerator(
+        data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
             preprocessing_function=preprocess_input)
         evaluation_generator = data_generator.flow_from_dataframe(
             self.evaluation_labels,
@@ -265,7 +268,7 @@ class EyeballModel:
         if not self.random_seed:
             print("Using validation set...")
             # Data augmentation
-            data_generator = ImageDataGenerator(
+            data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
                 preprocessing_function=self.preprocess_training_function,
                 samplewise_center=True,
                 validation_split=0.2)

eyeballer/visualization.py

@@ -1,17 +1,18 @@
 import numpy as np
 import math
-from keras.preprocessing import image
+import tensorflow as tf
+
 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):
+    def __init__(self, screenshot_path, model, threshold=0.5, 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)
+        img = tf.keras.preprocessing.image.load_img(screenshot_path, target_size=(self.model.image_width, self.model.image_height))
+        img = tf.keras.preprocessing.image.img_to_array(img)
 
         self.screenshot = img
         self.boxsize = boxsize
@@ -20,6 +21,7 @@ class HeatMap():
         self.y = 0
         self.screenshot_path = screenshot_path
         self.gamma = 3
+        self.threshold = threshold
 
     def generate(self, output_file="heatmap.png"):
         """ Make a single heatmap image and return it """
@@ -30,11 +32,11 @@ class HeatMap():
         for label in DATA_LABELS:
             boxsize = self.boxsize
             # Ignore this label if it didn't predict positively (true label)
-            if results[label] > 0.5:
+            if results[label] > self.threshold:
                 worst_score = 1
-                while worst_score > 0.5:
+                while worst_score > self.threshold:
                     heatmap, worst_score = self._get_heatmap(label, boxsize)
-                    if worst_score > 0.5:
+                    if worst_score > self.threshold:
                         boxsize += 28
                         print("Didn't get a good image for {}. Trying again with a bigger boxsize: {}".format(label, boxsize))
                     else:
@@ -42,7 +44,7 @@ class HeatMap():
                         labels.append(label)
 
         plt.figure()
-        screenshot_image = image.load_img(self.screenshot_path, target_size=(self.model.image_width, self.model.image_width))
+        screenshot_image = tf.keras.preprocessing.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

requirements-gpu.txt

@@ -5,5 +5,6 @@ numpy
 pandas
 pillow
 sklearn
-tensorflow-gpu
-keras
+tensorflow-gpu==2.0.0-beta0
+jinja2
+progressbar2

requirements.txt

@@ -1,10 +1,11 @@
 Augmentor
 click
 matplotlib
-numpy
+numpy==1.16.4
+keras
 pandas
 pillow
 sklearn
-tensorflow
-keras
+tensorflow==2.0.0-beta0
 jinja2
+progressbar2

tests/EvaluateTest.py

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

tests/PredictTest.py

@@ -1,15 +1,22 @@
 import errno
 import os
 import unittest
+import sys
 
 from eyeballer.model import EyeballModel
 
 
 class PredictTest(unittest.TestCase):
     def setUp(self):
-        weights_file = "latest_weights.h5"
+
+        class DummyFile(object):
+            def write(self, x): pass
+            def flush(self): pass
+
+        sys.stdout = DummyFile()
+        weights_file = "tests/models/test_weights.h5"
         if not os.path.isfile(weights_file):
-            print("Error: Symlink the latest weights file to 'latest_weights.h5'")
+            print("Error: Symlink the latest weights file to " + weights_file)
             raise FileNotFoundError(
                 errno.ENOENT,
                 os.strerror(errno.ENOENT),
@@ -18,7 +25,8 @@ class PredictTest(unittest.TestCase):
         model_kwargs = {
             "weights_file": weights_file,
             "print_summary": False,
-            "seed": None
+            "seed": None,
+            "quiet": True
         }
         self.model = EyeballModel(**model_kwargs)
 
@@ -26,11 +34,12 @@ class PredictTest(unittest.TestCase):
         model_kwargs = {
             "weights_file": None,
             "print_summary": False,
-            "seed": 12345678
+            "seed": 12345678,
+            "quiet": True
         }
         same_seed_model = EyeballModel(**model_kwargs)
 
-        screenshot = "tests/data/test.png"
+        screenshot = "tests/data/404.png"
 
         results_one = self.model.predict(screenshot)
         results_two = same_seed_model.predict(screenshot)
@@ -41,44 +50,82 @@ class PredictTest(unittest.TestCase):
         model_kwargs = {
             "weights_file": None,
             "print_summary": False,
-            "seed": 12345678
+            "seed": 12345678,
+            "quiet": True
         }
         same_seed_model = EyeballModel(**model_kwargs)
 
-        screenshot = "tests/data/test.png"
+        screenshot = "tests/data/404.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
+        self.assertEqual(results_one, results_two)
 
     def test_predict_custom404(self):
-        screenshot = "tests/data/custom404.png"
-        results = self.model.predict(screenshot)
-        self.assertEqual(self.get_prediction_by_name(results), "custom404")
+        screenshot = "tests/data/404.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["custom404"], 0.5)
+
+    def test_predict_not_custom404(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["custom404"], 0.5)
 
     def test_predict_login(self):
         screenshot = "tests/data/login.png"
-        results = self.model.predict(screenshot)
-        self.assertEqual(self.get_prediction_by_name(results), "login")
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["login"], 0.5)
+
+    def test_predict_not_login(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        print(screenshot, results)
+        self.assertLess(results["login"], 0.5)
 
     def test_predict_homepage(self):
         screenshot = "tests/data/homepage.png"
-        results = self.model.predict(screenshot)
-        self.assertEqual(self.get_prediction_by_name(results), "homepage")
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["homepage"], 0.5)
+
+    def test_predict_not_homepage(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["homepage"], 0.5)
 
     def test_predict_oldlooking(self):
-        screenshot = "tests/data/oldlooking.png"
-        results = self.model.predict(screenshot)
-        self.assertEqual(self.get_prediction_by_name(results), "oldlooking")
+        screenshot = "tests/data/old-looking.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["oldlooking"], 0.5)
+
+    def test_predict_not_oldlooking(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["oldlooking"], 0.5)
+
+    def test_file_doesnt_exist(self):
+        screenshot = "tests/data/doesnotexist.png"
+        try:
+            self.model.predict(screenshot)[0]
+            self.fail("FileNotFoundError was expected but not found")
+        except FileNotFoundError:
+            pass
+
+    def test_folder(self):
+        screenshots = "tests/data/"
+        results = self.model.predict(screenshots)
+        self.assertEqual(len(results), 7)
+
+    def test_file_is_empty(self):
+        """
+        We're just testing that it doesn't crash, basically
+        """
+        screenshot = "tests/data/empty.png"
+        self.model.predict(screenshot)
+
+    def test_file_is_invalid(self):
+        """
+        We're just testing that it doesn't crash, basically
+        """
+        screenshot = "tests/data/invalid.png"
+        self.model.predict(screenshot)

tests/TrainTest.py

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

tests/data/invalid.png

@@ -0,0 +1 @@
+THIS TEXT IS NOT A VALID PNG