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

README.md

@@ -7,6 +7,7 @@ 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
 
@@ -41,17 +42,22 @@ Parked domains are websites that look real, but aren't valid attack surface. The
 
 ## 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 requirements file has you covered.
+**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**
 
@@ -59,15 +65,20 @@ For the latest pretrained weights, check out the releases here on GitHub.
 
 **Training Data** You can find our training data here:
 
-https://www.dropbox.com/s/rpylhiv2g0kokts/eyeballer-3.0.zip?dl=1
+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
-3. `bishop-fox-pretrained-v3.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.
+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
 
@@ -76,13 +87,13 @@ NOTE: For best results, make sure you screenshot your websites in a native 1.6x
 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).
@@ -112,7 +123,7 @@ For a detailed explanation on [Precision vs Recall, check out Wikipedia](https:/
 ## 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.
@@ -124,7 +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/visualization.py

@@ -11,8 +11,11 @@ class HeatMap():
     def __init__(self, screenshot_path, model, threshold=0.5, boxsize=28, step=7):
         self.model = model
 
-        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)
+        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

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
-jinja2
-progressbar2

requirements.txt

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

tests/PredictTest.py

@@ -1,142 +0,0 @@
-import errno
-import os
-import unittest
-import sys
-
-from eyeballer.model import EyeballModel
-from eyeballer.visualization import HeatMap
-
-
-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,
-            "quiet": True
-        }
-        self.model = EyeballModel(**model_kwargs)
-
-    def test_different_seed_predict(self):
-        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 = 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,
-            "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)
-
-        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)
-
-    def test_heatmap(self):
-        screenshot = "tests/data/login.png"
-        HeatMap(screenshot, self.model, 0.5)
-
-        screenshot = "tests/data/404.png"
-        HeatMap(screenshot, self.model, 0.5)
-
-        screenshot = "tests/data/empty.png"
-        HeatMap(screenshot, self.model, 0.5)

tests/__init__.py

@@ -1 +1 @@
-from tests.PredictTest import *  # noqa
+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/verify.py

@@ -15,11 +15,14 @@ with open("labels.csv", newline="") as csvfile:
     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(row["filename"])
+            print(label)
 
     print("Images that exist, but don't have labels:")
     i = 0