1 files changed, 16 insertions, 13 deletions
diff --git a/paste/__main__.py b/paste/__main__.py
index 27ff72b..6e684c4 100644
--- a/paste/__main__.py
+++ b/paste/__main__.py
@@ -1,11 +1,12 @@
 from base64 import b64decode, b64encode
-from contextlib import AbstractContextManager
+from collections.abc import Iterator
+from contextlib import contextmanager
 from datetime import datetime, timezone
 from os import getenv
 from sys import argv, stderr
 from wsgiref.simple_server import make_server
 
-from . import DB_PATH, application, store
+from . import DB_PATH, application, db, store
 
 PROGRAM_NAME = "paste"
 
@@ -42,8 +43,10 @@ Environment:
 db_path = getenv("PASTE_DB", DB_PATH)
 
 
-def open_db() -> AbstractContextManager[store.Store]:
-    return store.open(db_path)
+@contextmanager
+def open_auth() -> Iterator[store.Auth]:
+    with db.connect(db_path) as conn:
+        yield store.Auth(conn)
 
 
 def main():
@@ -54,11 +57,11 @@ def main():
         httpd.serve_forever()
     elif len(argv) == 2:
         if argv[1] == "new-token":
-            with open_db() as db:
-                print(b64encode(db.generate_token()).decode())
+            with open_auth() as auth:
+                print(b64encode(auth.generate_token()).decode())
         elif argv[1] == "list-tokens":
-            with open_db() as db:
-                for token_hash, created_at in db.get_tokens():
+            with open_auth() as auth:
+                for token_hash, created_at in auth.get_tokens():
                     created_at = datetime.fromtimestamp(created_at, timezone.utc)
                     print(f"{token_hash.hex()}\t{created_at.ctime()}")
         elif argv[1] == "-h" or argv[1] == "--help":
@@ -70,12 +73,12 @@ def main():
     elif len(argv) == 3:
         if argv[1] == "delete-token":
             token = argv[2]
-            with open_db() as db:
+            with open_auth() as auth:
                 try:
                     try:
-                        db.delete_token(b64decode(token))
+                        auth.delete_token(b64decode(token))
                     except ValueError:
-                        db.delete_token_hash(bytes.fromhex(token))
+                        auth.delete_token_hash(bytes.fromhex(token))
                 except ValueError:
                     print("Malformed token", file=stderr)
                     exit(1)
@@ -83,9 +86,9 @@ def main():
                     print("Token not found", file=stderr)
                     exit(1)
         elif argv[1] == "verify-token":
-            with open_db() as db:
+            with open_auth() as auth:
                 try:
-                    if not db.check_token(b64decode(argv[2])):
+                    if not auth.check_token(b64decode(argv[2])):
                         print("Token not found", file=stderr)
                         exit(1)
                     print("Found")