Unit tests in github workflow

- add uv
 - fix unit tests

.github/workflows/tests.yml

@@ -0,0 +1,21 @@
+name: Tests
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Install dependencies
+        run: uv sync --group dev
+
+      - name: Run tests
+        run: uv run pytest tests/ -v

.gitignore

@@ -118,3 +118,5 @@ results.html
 !requirements-gpu.txt
 !tests/models/*
 !tests/data/*
+!docs/*
+logs/*

README.md

@@ -1,37 +1,129 @@
 # Eyeballer
+
+![Logo](/docs/eyeballer_logo.png)
+
+
 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.
 
+Give it a try live at: https://eyeballer.bishopfox.com
+
+### Example Labels
+
+| Old-Looking Sites | Login Pages |
+| ------ |:-----:|
+| ![Sample Old-looking Page](/docs/old-looking.png) | ![Sample Login Page](/docs/login.png) |
+
+| Webapp | Custom 404's |
+| ------ |:-----:|
+| ![Sample Webapp](/docs/homepage.png) | ![Sample Custom 404](/docs/404.png) |
+
+| Parked Domains |
+| ------ |
+| ![Sample Webapp](/docs/parked.png) |
+
+## What the Labels Mean
+
+**Old-Looking Sites**
+Blocky frames, broken CSS, that certain "je ne sais quoi" of a website that looks like it was designed in the early 2000's. You know it when you see it. Old websites aren't just ugly, they're also typically super vulnerable. When you're looking to hack into something, these websites are a gold mine.
+
+**Login Pages**
+Login pages are valuable to pen testing, they indicate that there's additional functionality you don't currently have access to. It also means there's a simple follow-up process of credential enumeration attacks. You might think that you can set a simple heuristic to find login pages, but in practice it's really hard. Modern sites don't just use a simple input tag we can grep for.
+
+**Webapp**
+This tells you that there is a larger group of pages and functionality available here that can serve as surface area to attack. This is in contrast to a simple login page, with no other functionality. Or a default IIS landing page which has no other functionality. This label should indicate to you that there is a web application here to attack.
+
+**Custom 404**
+Modern sites love to have cutesy custom 404 pages with pictures of broken robots or sad looking dogs. Unfortunately, they also love to return HTTP 200 response codes while they do it. More often, the "404" page doesn't even contain the text "404" in it. These pages are typically uninteresting, despite having a lot going on visually, and Eyeballer can help you sift them out.
+
+**Parked Domains**
+Parked domains are websites that look real, but aren't valid attack surface. They're stand-in pages, usually devoid of any real functionality, consist almost entirely of ads, and are usually not run by our actual target. It's what you get when the domain specified is wrong or lapsed. Finding these pages and removing them from scope is really valuable over time.
+
 ## Setup
 
-Download required packages on pip:
+Install [uv](https://docs.astral.sh/uv/getting-started/installation/) if you haven't already:
 ```
-sudo pip3 install -r requirements.txt
+curl -LsSf https://astral.sh/uv/install.sh | sh
+```
+
+Then install eyeballer and its dependencies:
+```
+uv sync
 ```
 
 Or if you want GPU support:
 ```
-sudo pip3 install -r requirements-gpu.txt
+uv sync --extra gpu
 ```
 
+**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 command has you covered.
+
+**Pretrained Weights**
+
+For the latest pretrained weights, check out the releases here on GitHub.
+
+**Training Data** You can find our training data here:
+
+https://www.kaggle.com/altf42600/pentest-screensots
+
+There's two things you need from the training data:
+
+1. `images/` folder, containing all the screenshots (resized down to 224x224)
+2. `labels.csv` that has all the labels
+
+Copy both into the root of the Eyeballer code tree.
+
+Additionally, you can find a pretrained weights file you can use right out of the box without training.  
+
+1. `bishop-fox-pretrained-vN.h5`
+
+It's here on GitHub, look at the `releases` section for the latest one.
+
 ## Predicting Labels
+
+NOTE: For best results, make sure you screenshot your websites in a native 1.6x aspect ratio. IE: 1440x900. Eyeballer will scale the image down automatically to the right size for you, but if it's the wrong aspect ratio then it will squish in a way that will affect prediction performance.
+
 To eyeball some screenshots, just run the "predict" mode:
 
 ```
-eyeballer.py --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
+uv run eyeballer.py --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
 ```
 
 Or for a whole directory of files:
 
 ```
-eyeballer.py --weights YOUR_WEIGHTS.h5 predict PATH_TO/YOUR_FILES/
+uv run 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).
+
+## Performance
+
+Eyeballer's performance is measured against an evaluation dataset, which is 20% of the overall screenshots chosen at random. Since these screenshots are never used in training, they can be an effective way to see how well the model is performing. Here are the latest results:
+
+| Overall Binary Accuracy | 93.52% |
+| ------ |:-----:|
+| **All-or-Nothing Accuracy** | **76.09%** |
+
+**Overall Binary Accuracy** is probably what you think of as the model's "accuracy". It's the chance, given any single label, that it is correct.
+
+**All-or-Nothing Accuracy** is more strict. For this, we consider all of an image's labels and consider it a failure if ANY label is wrong. This accuracy rating is the chance that the model correctly predicts all labels for any given image.
+
+| Label | Precision | Recall |
+| ------ | ------ |:-----:|
+| Custom 404 | 80.20% | 91.01% |
+| Login Page | 86.41% | 88.47% |
+| Webapp | 95.32% | 96.83% |
+| Old Looking | 91.70% | 62.20% |
+| Parked Domain | 70.99% | 66.43% |
+
+For a detailed explanation on [Precision vs Recall, check out Wikipedia](https://en.wikipedia.org/wiki/Precision_and_recall).
+
 ## Training
 To train a new model, run:
 ```
-eyeballer.py train
+uv run 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 +135,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.py --weights YOUR_WEIGHTS.h5 evaluate
+uv run 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

@@ -54,7 +54,7 @@ def predict(ctx, screenshot, heatmap, threshold):
         print(results)
     else:
         with open("results.csv", "w", newline="") as csvfile:
-            fieldnames = ["filename", "custom404", "login", "homepage", "oldlooking"]
+            fieldnames = ["filename", "custom404", "login", "webapp", "oldlooking", "parked"]
             labelwriter = csv.DictWriter(csvfile, fieldnames=fieldnames)
             labelwriter.writeheader()
             labelwriter.writerows(results)

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

@@ -17,7 +17,7 @@ from tensorflow.keras.applications.mobilenet import preprocess_input
 from sklearn.metrics import classification_report, accuracy_score, hamming_loss
 from eyeballer.augmentation import EyeballerAugmentation
 
-DATA_LABELS = ["custom404", "login", "homepage", "oldlooking"]
+DATA_LABELS = ["custom404", "login", "webapp", "oldlooking", "parked"]
 
 
 class EyeballModel:
@@ -31,7 +31,7 @@ class EyeballModel:
     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:
@@ -54,28 +54,6 @@ class EyeballModel:
                            loss="binary_crossentropy",
                            metrics=["accuracy"])
 
-        # for layer in pretrained_layer.layers:
-        #     layer.trainable = False
-
-        if print_summary:
-            print(self.model.summary())
-
-        # 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
-        self.evaluation_labels = data.loc[data['evaluation'] == True]  # noqa: E712
-
-        # Shuffle the training labels
-        self.random_seed = False
-        self.seed = seed
-        if self.seed is None:
-            self.random_seed = True
-            self.seed = random.randint(0, 999999)
-            print("No seed set, ", end='')
-        print("using seed: {}".format(self.seed))
-        random.seed(self.seed)
-        self.training_labels = self.training_labels.sample(frac=1)
-
         if weights_file is not None and os.path.isfile(weights_file):
             try:
                 self.model.load_weights(weights_file)
@@ -88,6 +66,28 @@ class EyeballModel:
                 raise FileNotFoundError
             print("WARN: No model loaded from file. Generating random model")
 
+        if print_summary:
+            print(self.model.summary())
+
+        self.quiet = quiet
+        self.seed = seed
+
+    def _init_labels(self):
+        # 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
+        self.evaluation_labels = data.loc[data['evaluation'] == True]  # noqa: E712
+
+        # Shuffle the training labels
+        self.random_seed = False
+        if self.seed is None:
+            self.random_seed = True
+            self.seed = random.randint(0, 999999)
+            print("No seed set, ", end='')
+        print("using seed: {}".format(self.seed))
+        random.seed(self.seed)
+        self.training_labels = self.training_labels.sample(frac=1)
+
         # Data augmentation
         augmentor = Augmentor.Pipeline()
         augmentor.set_seed(self.seed)
@@ -112,6 +112,7 @@ class EyeballModel:
         """
         print("Training with seed: " + str(self.seed))
 
+        self._init_labels()
         data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
             preprocessing_function=self.preprocess_training_function,
             validation_split=0.2,
@@ -127,7 +128,7 @@ class EyeballModel:
             subset='training',
             shuffle=True,
             seed=self.seed,
-            class_mode="other")
+            class_mode="raw")
         validation_generator = data_generator.flow_from_dataframe(
             self.training_labels,
             directory=self.image_dir,
@@ -138,17 +139,26 @@ class EyeballModel:
             subset='validation',
             shuffle=False,
             seed=self.seed,
-            class_mode="other")
+            class_mode="raw")
 
         # Model checkpoint - Saves model weights when validation accuracy improves
         callbacks = [tf.keras.callbacks.ModelCheckpoint(self.checkpoint_file,
-                     monitor='val_loss',
-                     verbose=1,
-                     save_best_only=True,
-                     save_weights_only=True,
-                     mode='min')]
+                                                        monitor='val_loss',
+                                                        verbose=1,
+                                                        save_best_only=True,
+                                                        save_weights_only=True,
+                                                        mode='min'),
+                     tf.keras.callbacks.TensorBoard(log_dir='logs',
+                                                    histogram_freq=2,
+                                                    write_graph=True,
+                                                    write_images=False,
+                                                    update_freq='epoch',
+                                                    profile_batch=2,
+                                                    embeddings_freq=0,
+                                                    embeddings_metadata=None),
+                     ]
 
-        history = self.model.fit_generator(
+        history = self.model.fit(
             training_generator,
             steps_per_epoch=len(training_generator.filenames) // batch_size,
             epochs=epochs,
@@ -195,8 +205,9 @@ class EyeballModel:
         result["filename"] = "custom-image"
         result["custom404"] = prediction[0][0]
         result["login"] = prediction[0][1]
-        result["homepage"] = prediction[0][2]
+        result["webapp"] = prediction[0][2]
         result["oldlooking"] = prediction[0][3]
+        result["parked"] = prediction[0][4]
         return result
 
     def predict(self, path, threshold=0.5):
@@ -216,7 +227,10 @@ class EyeballModel:
             raise FileNotFoundError
 
         results = []
-        for screenshot in progressbar.progressbar(screenshots):
+        bar = progressbar.ProgressBar()
+        if self.quiet:
+            bar = progressbar.NullBar()
+        for screenshot in bar(screenshots):
             # Load the image into memory
             img = None
             try:
@@ -236,8 +250,9 @@ class EyeballModel:
             result["filename"] = screenshot
             result["custom404"] = prediction[0][0]
             result["login"] = prediction[0][1]
-            result["homepage"] = prediction[0][2]
+            result["webapp"] = prediction[0][2]
             result["oldlooking"] = prediction[0][3]
+            result["parked"] = prediction[0][4]
             results.append(result)
         return results
 
@@ -248,6 +263,7 @@ class EyeballModel:
         threshold -- Value between 0->1. The cutoff where the numerical prediction becomes boolean. Default: 0.5
         """
         # Prepare the data
+        self._init_labels()
         data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
             preprocessing_function=preprocess_input)
         evaluation_generator = data_generator.flow_from_dataframe(
@@ -258,7 +274,7 @@ class EyeballModel:
             target_size=(self.image_width, self.image_height),
             shuffle=False,
             batch_size=1,
-            class_mode="other")
+            class_mode="raw")
         # If a seed was selected, then also evaluate on the validation set for that seed
         if not self.random_seed:
             print("Using validation set...")
@@ -277,17 +293,16 @@ class EyeballModel:
                 subset='validation',
                 shuffle=False,
                 seed=self.seed,
-                class_mode="other")
+                class_mode="raw")
         else:
             print("Using evaluation set...")
 
-        predictions = self.model.predict_generator(
+        predictions = self.model.predict(
             evaluation_generator,
             verbose=1,
             steps=len(evaluation_generator))
 
         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()

eyeballer/visualization.py

@@ -1,6 +1,7 @@
 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
@@ -10,8 +11,11 @@ class HeatMap():
     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)
+        try:
+            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)
+        except Exception:
+            img = None
 
         self.screenshot = img
         self.boxsize = boxsize
@@ -43,7 +47,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

prediction_output_template.html

@@ -7,8 +7,9 @@
       <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('webapp')"> Web Application</button>
       <button class="btn" onclick="filterSelection('oldlooking')"> Old Looking</button>
+      <button class="btn" onclick="filterSelection('parked')"> Parked Domain</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>

pyproject.toml

@@ -0,0 +1,35 @@
+[project]
+name = "eyeballer"
+version = "3.0.0"
+description = "Give those screenshots of yours a quick eyeballing."
+requires-python = ">=3.8"
+dependencies = [
+    "Augmentor",
+    "click",
+    "matplotlib",
+    "numpy",
+    "pandas",
+    "pillow",
+    "scikit-learn",
+    "tensorflow",
+    "jinja2",
+    "progressbar2",
+    "scipy",
+]
+
+[project.optional-dependencies]
+gpu = [
+    "tensorflow[and-cuda]",
+]
+
+[project.scripts]
+eyeballer = "eyeballer:cli"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3.5",
+]

requirements-gpu.txt

@@ -1,10 +0,0 @@
-Augmentor
-click
-matplotlib
-numpy
-pandas
-pillow
-sklearn
-tensorflow-gpu==2.0.0-beta0
-jinja2
-progressbar2

requirements.txt

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

tests/EvaluateTest.py

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

tests/PredictTest.py

@@ -1,128 +0,0 @@
-import errno
-import os
-import unittest
-import sys
-
-from eyeballer.model import EyeballModel
-
-
-class PredictTest(unittest.TestCase):
-    def setUp(self):
-
-        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 " + weights_file)
-            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/404.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/404.png"
-
-        results_one = same_seed_model.predict(screenshot)
-        results_two = same_seed_model.predict(screenshot)
-
-        self.assertEqual(results_one, results_two)
-
-    def test_predict_custom404(self):
-        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)[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)[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/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/__init__.py

@@ -1 +1 @@
-from tests.PredictTest import *
+from tests.test_predict import *  # noqa

tests/test_predict.py

@@ -0,0 +1,141 @@
+import os
+import pytest
+
+from eyeballer.model import EyeballModel
+from eyeballer.visualization import HeatMap
+
+
+@pytest.fixture
+def model():
+    weights_file = "tests/models/test_weights.h5"
+    if not os.path.isfile(weights_file):
+        pytest.skip("Symlink the latest weights file to " + weights_file)
+
+    model_kwargs = {
+        "weights_file": weights_file,
+        "print_summary": False,
+        "seed": None,
+        "quiet": True
+    }
+    return EyeballModel(**model_kwargs)
+
+
+def test_different_seed_predict(model):
+    model_kwargs = {
+        "weights_file": None,
+        "print_summary": False,
+        "seed": 12345678,
+        "quiet": True
+    }
+    same_seed_model = EyeballModel(**model_kwargs)
+
+    screenshot = "tests/data/404.png"
+
+    results_one = model.predict(screenshot)
+    results_two = same_seed_model.predict(screenshot)
+
+    assert results_one != results_two
+
+
+def test_same_seed_predict():
+    model_kwargs = {
+        "weights_file": None,
+        "print_summary": False,
+        "seed": 12345678,
+        "quiet": True
+    }
+    same_seed_model = EyeballModel(**model_kwargs)
+
+    screenshot = "tests/data/404.png"
+
+    results_one = same_seed_model.predict(screenshot)
+    results_two = same_seed_model.predict(screenshot)
+
+    assert results_one == results_two
+
+
+def test_predict_custom404(model):
+    screenshot = "tests/data/404.png"
+    results = model.predict(screenshot)[0]
+    assert results["custom404"] > 0.5
+
+
+def test_predict_not_custom404(model):
+    screenshot = "tests/data/nothing.png"
+    results = model.predict(screenshot)[0]
+    assert results["custom404"] < 0.5
+
+
+def test_predict_login(model):
+    screenshot = "tests/data/login.png"
+    results = model.predict(screenshot)[0]
+    assert results["login"] > 0.5
+
+
+def test_predict_not_login(model):
+    screenshot = "tests/data/nothing.png"
+    results = model.predict(screenshot)[0]
+    assert results["login"] < 0.5
+
+
+def test_predict_homepage(model):
+    screenshot = "tests/data/homepage.png"
+    results = model.predict(screenshot)[0]
+    assert results["webapp"] > 0.5
+
+
+def test_predict_not_homepage(model):
+    screenshot = "tests/data/nothing.png"
+    results = model.predict(screenshot)[0]
+    assert results["webapp"] < 0.5
+
+
+def test_predict_oldlooking(model):
+    screenshot = "tests/data/old-looking.png"
+    results = model.predict(screenshot)[0]
+    assert results["oldlooking"] > 0.5
+
+
+def test_predict_not_oldlooking(model):
+    screenshot = "tests/data/nothing.png"
+    results = model.predict(screenshot)[0]
+    assert results["oldlooking"] < 0.5
+
+
+def test_file_doesnt_exist(model):
+    screenshot = "tests/data/doesnotexist.png"
+    with pytest.raises(FileNotFoundError):
+        model.predict(screenshot)
+
+
+def test_folder(model):
+    screenshots = "tests/data/"
+    results = model.predict(screenshots)
+    assert len(results) == 7
+
+
+def test_file_is_empty(model):
+    """
+    We're just testing that it doesn't crash, basically
+    """
+    screenshot = "tests/data/empty.png"
+    model.predict(screenshot)
+
+
+def test_file_is_invalid(model):
+    """
+    We're just testing that it doesn't crash, basically
+    """
+    screenshot = "tests/data/invalid.png"
+    model.predict(screenshot)
+
+
+def test_heatmap(model):
+    screenshot = "tests/data/login.png"
+    HeatMap(screenshot, model, 0.5)
+
+    screenshot = "tests/data/404.png"
+    HeatMap(screenshot, model, 0.5)
+
+    screenshot = "tests/data/empty.png"
+    HeatMap(screenshot, model, 0.5)

utils/resizer.sh

@@ -1,7 +1,8 @@
 #!/bin/bash
 
 # Converts full-size images to a given size (replacing the originals)
-for f in images_full_size/*
+for f in images/*
 do
-	convert -resize 224x224 $f $f
+	echo $f
+	convert -resize 224x224\! $f $f
 done

utils/verify.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+
+import csv
+import random
+from os import listdir
+from os.path import isfile, join
+
+with open("labels.csv", newline="") as csvfile:
+    # Get the header labels
+    csvreader = csv.DictReader(csvfile)
+    # fieldnames = next(csvreader)
+    file_list = [f for f in listdir("images") if isfile(join("images", f))]
+
+    label_list = []
+    for row in csvreader:
+        label_list.append(row["filename"])
+
+    print(len(file_list), "files")
+    print(len(label_list), "labels")
+
+    # Loop through the file
+    print("Rows in labels.csv, but no image exists:")
+    for label in label_list:
+        if label not in file_list:
+            print(label)
+
+    print("Images that exist, but don't have labels:")
+    i = 0
+    for image in file_list:
+        if image not in label_list:
+            print(image)