GH-4 # add check_password_leak logic

Author: Hugo-C

poetry.lock

@@ -26,6 +26,7 @@ redis = "^5.2.1"
 docopt = "^0.6.2"
 xxhash = "^3.5.0"
 fastapi = {extras = ["standard"], version = "^0.115.8"}
+pydantic-settings = "^2.7.1"
 
 [tool.ruff]
 line-length = 120

pyproject.toml

@@ -1,3 +1,4 @@
+from functools import cache
 from typing import Annotated
 
 from fastapi.params import Depends
@@ -9,6 +10,11 @@ def get_settings() -> Settings:
     return Settings()
 
 
-# TODO cache settings + password_storage ?
+@cache
+def get_settings_cached() -> Settings:
+    return Settings()
+
+
+# TODO cache password_storage ?
 
-SettingsDep = Annotated[Settings, Depends(get_settings)]
+SettingsDep = Annotated[Settings, Depends(get_settings_cached)]

src/api/dependencies.py

@@ -1,6 +1,7 @@
 from typing import Annotated
 
 from fastapi import APIRouter, Path
+from redis import Redis
 
 from src.api.dependencies import SettingsDep
 from src.common import PasswordStorage
@@ -20,6 +21,11 @@
 ]
 
 
-@v1_router.get("/haveibeenrocked/{prefix}")
+@v1_router.get("/haveibeenrocked/{prefix}")  # TODO add response schema to openAPI
 def check_password_leak(prefix: APIPrefix, settings: SettingsDep):
-    return {"Hello": "World"}
+    redis_client = Redis.from_url(str(settings.kvrocks_url))
+    password_storage = PasswordStorage(client=redis_client)
+
+    matching_password = password_storage.get_passwords(prefix=prefix)
+
+    return {prefix: matching_password}

src/api/router.py

@@ -1,5 +1,5 @@
 from pydantic import RedisDsn
-from pydantic.v1 import BaseSettings
+from pydantic_settings import BaseSettings
 
 
 class Settings(BaseSettings):

src/common/settings.py

@@ -1,19 +1,36 @@
+import os
+from unittest import mock
+
 import pytest
 from redis import Redis
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
+from src.api.dependencies import get_settings, get_settings_cached
+from src.api.main import app
+from src.common import PasswordStorage
+
 KVROCKS_IMAGE = "apache/kvrocks:2.11.0"
 
 
+@pytest.fixture(scope="session", autouse=True)
+def settings():  # disable the cache on settings, so it'll be reloaded everytime
+    app.dependency_overrides[get_settings_cached] = get_settings
+    try:
+        yield get_settings  # to give tests access to settings easily if needed
+    finally:
+        app.dependency_overrides = {}
+
+
 @pytest.fixture(scope="session")
 def _kvrocks() -> Redis:
     with DockerContainer(image=KVROCKS_IMAGE).with_exposed_ports(6666) as kvrocks_container:
         wait_for_logs(kvrocks_container, "Ready to accept connections")
         host = kvrocks_container.get_container_host_ip()
         port = kvrocks_container.get_exposed_port(6666)
 
-        yield Redis(host=host, port=port)
+        with mock.patch.dict(os.environ, {"KVROCKS_URL": f"redis://{host}:{port}"}):
+            yield Redis(host=host, port=port)
 
 
 @pytest.fixture()
@@ -24,3 +41,10 @@ def kvrocks(_kvrocks) -> Redis:
         # Clean up keys after the test is done
         for key in _kvrocks.keys():
             _kvrocks.delete(key)
+
+
+@pytest.fixture
+def password_storage(kvrocks):
+    password_storage_ = PasswordStorage(client=kvrocks)
+    password_storage_.PIPELINE_MAX_SIZE = 1  # so we don't have to flush in our test
+    yield password_storage_

tests/conftest.py

@@ -0,0 +1,41 @@
+from fastapi.testclient import TestClient
+
+from src.api.main import app
+
+client = TestClient(app)
+
+PREFIX_CHECKED = "123AB"
+
+
+def test_api_return_no_results(kvrocks):
+    response = client.get(f"/api/v1/haveibeenrocked/{PREFIX_CHECKED}")
+
+    assert response.status_code == 200
+    assert response.json() == {PREFIX_CHECKED: []}
+
+
+def test_api_return_one_result(kvrocks, password_storage):
+    expected_password = "some password"
+    password_storage.add_password(prefix=PREFIX_CHECKED, password=expected_password)
+
+    response = client.get(f"/api/v1/haveibeenrocked/{PREFIX_CHECKED}")
+
+    assert response.status_code == 200
+    assert response.json() == {PREFIX_CHECKED: [expected_password]}
+
+
+def test_api_return_multiple_results(kvrocks, password_storage):
+    expected_passwords = {
+        "some password 1",
+        "some password 2",
+        "some password 3",
+    }
+    for password in expected_passwords:
+        password_storage.add_password(prefix=PREFIX_CHECKED, password=password)
+
+    response = client.get(f"/api/v1/haveibeenrocked/{PREFIX_CHECKED}")
+
+    assert response.status_code == 200
+    json_response = response.json()
+    assert set(json_response.pop(PREFIX_CHECKED)) == expected_passwords
+    assert len(json_response) == 0  # no other keys

tests/integration/api/test_password_leak.py

@@ -5,13 +5,6 @@
 SOME_PREFIX = "01234A"
 
 
-@pytest.fixture
-def password_storage(kvrocks):
-    password_storage_ = PasswordStorage(client=kvrocks)
-    password_storage_.PIPELINE_MAX_SIZE = 1  # so we don't have to flush in our test
-    yield password_storage_
-
-
 def test_get_password_should_return_empty_set_at_first(password_storage):
     assert password_storage.get_passwords(SOME_PREFIX) == set()