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

@@ -42,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**
 
@@ -82,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).
@@ -118,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.
@@ -130,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,11 +0,0 @@
-Augmentor
-click
-matplotlib
-numpy
-pandas
-pillow
-scikit-learn
-tensorflow
-jinja2
-progressbar2
-scipy

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)