GH-6 # clean up init_db script and add download option

Author: Hugo-C

.dockerignore

@@ -1,4 +1,14 @@
 src/static/*.bundle.js
 package-lock.json
 
-node_modules/
+node_modules/
+
+.git/
+.idea/
+.venv/
+
+__pycache__/
+**/__pycache__/
+
+*.py[cod]
+*$py.class

README.md

@@ -91,13 +91,15 @@ docker compose up
 
 Dans un nouveau terminal, si vous avez déjà une base de mot de passe type rockyou.txt en local:
 ```shell
-docker compose run -v <path_to_directory_having_rockyou.txt>:/rockdir fastapi python src/init_scripts/init_db_with_password_file.py /rockdir/rockyou.txt
+docker compose run --rm -v <path_to_directory_having_rockyou.txt>:/rockdir fastapi python src/init_scripts/init_db_with_password_file.py /rockdir/rockyou.txt
 ```
 si vous souhaitez le télécharger automatiquement:
 ```shell
-TODO
+docker compose run --rm fastapi python src/init_scripts/init_db_with_password_file.py --download
 ```
 
+La base de donnée utilise un volume Docker qui peut être supprimé avec `docker volume remove hibp_kvrocks_data_volume`.
+
 ### Setup dev
 
 `poetry sync` permet d'installer les dépendances. L'outils justfile (voir [ici pour l'installation](https://github.com/casey/just?tab=readme-ov-file#cross-platform)) permet de faire tourner les tests et le linter.  

poetry.lock

@@ -18,7 +18,6 @@ build-backend = "poetry.core.masonry.api"
 pytest = "^8.3.4"
 ruff = "^0.9.4"
 testcontainers = {extras = ["generic"], version = "^4.9.1"}
-httpx = "^0.28.1"  # required to test fastapi
 
 
 [tool.poetry.dependencies]
@@ -27,6 +26,8 @@ docopt = "^0.6.2"
 xxhash = "^3.5.0"
 fastapi = {extras = ["standard"], version = "^0.115.8"}
 pydantic-settings = "^2.7.1"
+httpx = "^0.28.1"
+rich = "^13.9.4"
 
 [tool.ruff]
 line-length = 120

pyproject.toml

@@ -12,45 +12,76 @@
 
 """
 
-import os
+import tempfile
 import time
 
+import httpx
+import rich
 from docopt import docopt
 from redis import Redis
+from rich import print
+from rich.progress import MofNCompleteColumn, Progress, SpinnerColumn
 
 from src.common import PasswordHasher, PasswordStorage, Settings
 
+ROCKYOU_DOWNLOAD_URL = "https://github.com/brannondorsey/naive-hashcat/releases/download/data/rockyou.txt"
 
-def init_db(file_path: str, db_client: Redis) -> int:
+
+def init_db(file_path: str, db_client: Redis, password_count: int) -> int:
     """Return the number of password inserted in db"""
-    print("")  # allow print below to clear this line return
+    processed = 0
     hasher = PasswordHasher()
     storage = PasswordStorage(client=db_client)
-    processed = 0
-    with open(file_path, encoding="latin-1") as file:
-        for line in file:
-            password = line.strip()
-            prefix = hasher.prefix(password)
-            storage.add_password(prefix=prefix, password=password)
-            processed += 1
-            if processed % 100 == 0:
-                print(f"\r{processed=}", end="")
+
+    with Progress(
+        SpinnerColumn(),
+        *Progress.get_default_columns(),
+        MofNCompleteColumn(),
+    ) as progress:
+        rich_task = progress.add_task("[green]Inserting in database...", total=password_count)
+        # Actual work
+        with open(file_path, encoding="latin-1") as file:
+            for line in file:
+                password = line.strip()
+                prefix = hasher.prefix(password)
+                storage.add_password(prefix=prefix, password=password)
+
+                processed += 1
+                progress.update(rich_task, advance=1)
+
     storage.flush()
-    print("\ndone")
+    print("[green]done")
     return processed
 
 
 if __name__ == "__main__":
+    start = time.time()
     args = docopt(__doc__)
     if args.get("--download"):
-        print("not implemented yet")
-        exit(0)
+        print("Initializing DB with rockyou passwords from internet")
+        # Create a temporary file that is deleted automatically on program close
+        download_file = tempfile.NamedTemporaryFile(delete=True, delete_on_close=False)
+        with httpx.stream("GET", ROCKYOU_DOWNLOAD_URL, follow_redirects=True) as r:
+            for data in r.iter_bytes():
+                download_file.write(data)
+        download_file.close()
+        print("download complete")
+        password_path = download_file.name
     elif password_path := args.get("<passwords-path>"):
         print(f"Initializing db with {password_path}")
-        start = time.time()
-        settings = Settings()
-        redis_client = Redis.from_url(str(settings.kvrocks_url))
-        init_db(password_path, db_client=redis_client)
-        duration = time.time() - start
-        print(f"Took {duration:.2f}s")
-    # else docopt will show help
+    else:
+        print("Missing arguments")
+        exit(1)
+
+    settings = Settings()
+    password_count = 0
+    with rich.progress.open(
+        password_path, encoding="latin-1", description="[orange3]Determining password count...", transient=True,
+    ) as file:
+        for line in file:
+            password_count += 1
+
+    redis_client = Redis.from_url(str(settings.kvrocks_url))
+    init_db(password_path, db_client=redis_client, password_count=password_count)
+    duration = time.time() - start
+    print(f"Took {duration:.2f}s")