From 7b3dc235f31bab7dcb0f9d37acfc86818791e697 Mon Sep 17 00:00:00 2001
From: Aedial <aedial.dev@gmail.com>
Date: Sun, 21 May 2023 17:44:22 +0200
Subject: [PATCH] [TOOL] Move decode into CLI

---
 README.md               |  6 ++++++
 README_TEMPLATE.md      |  6 ++++++
 novelai_api/__main__.py | 44 +++++++++++++++++++++++++++++++++++------
 tools/decode.py         | 17 ----------------
 4 files changed, 50 insertions(+), 23 deletions(-)
 delete mode 100644 tools/decode.py

diff --git a/README.md b/README.md
index 244c498..bc4fc21 100644
--- a/README.md
+++ b/README.md
@@ -39,6 +39,12 @@ Run a sanity check on your user content. It will print what content couldn't be
 python -m novelai_api sanity_check <username> <password>
 ```
 
+### Decode
+Decode a b64 encoded tokenized text. This will print the tokens and the decoded text.
+```bash
+python -m novelai_api decode <model> <data>
+```
+
 ## Using the module in your code
 A full list of examples is available in the [example](example) directory
 
diff --git a/README_TEMPLATE.md b/README_TEMPLATE.md
index 025e488..ca22b95 100644
--- a/README_TEMPLATE.md
+++ b/README_TEMPLATE.md
@@ -39,6 +39,12 @@ Run a sanity check on your user content. It will print what content couldn't be
 python -m novelai_api sanity_check <username> <password>
 ```
 
+### Decode
+Decode a b64 encoded tokenized text. This will print the tokens and the decoded text.
+```bash
+python -m novelai_api decode <model> <data>
+```
+
 ## Using the module in your code
 A full list of examples is available in the [example](example) directory
 
diff --git a/novelai_api/__main__.py b/novelai_api/__main__.py
index aeb0217..d5dbc8b 100644
--- a/novelai_api/__main__.py
+++ b/novelai_api/__main__.py
@@ -1,4 +1,7 @@
 import asyncio
+import base64
+import inspect
+import sys
 from argparse import ArgumentParser
 from logging import Logger, StreamHandler
 from typing import Any, Dict, List, NoReturn, Optional, Set, Tuple
@@ -6,6 +9,8 @@
 from aiohttp import ClientSession
 
 from novelai_api import NovelAIAPI
+from novelai_api.Preset import Model
+from novelai_api.Tokenizer import Tokenizer
 from novelai_api.utils import decompress_user_data, decrypt_user_data, get_access_key, get_encryption_key
 
 
@@ -158,32 +163,59 @@ async def sanity_checker_func(username: str, password: str):
         _check_duplicate_meta(stories, story_contents, presets, modules)
 
 
+# decode
+async def decode_func(model: str, data: str):
+    model = Model(model)
+
+    tokens = base64.b64decode(data)
+    tokens = [int.from_bytes(tokens[i * 2 : (i + 1) * 2], "little") for i in range(len(tokens) // 2)]
+    print(f"Tokens = {tokens}")
+
+    text = Tokenizer.decode(model, tokens)
+    print(f"Text = {text}")
+
+
 if __name__ == "__main__":
     parser = ArgumentParser()
 
+    def add_credentials_arguments(p: ArgumentParser):
+        p.add_argument("username", help="NovelAI username")
+        p.add_argument("password", help="NovelAI password")
+
     subparser = parser.add_subparsers(help="Function to call")
 
     # Get access key
     access_key_parser = subparser.add_parser("access_key", help="Get access key")
     access_key_parser.set_defaults(func=get_access_key_func)
+    add_credentials_arguments(access_key_parser)
 
     # Get access token
     access_token_parser = subparser.add_parser("access_token", help="Get access token")
     access_token_parser.set_defaults(func=get_access_token_func)
+    add_credentials_arguments(access_token_parser)
 
     # Sanity check
     sanity_check_parser = subparser.add_parser("sanity_check", help="Sanity check")
     sanity_check_parser.set_defaults(func=sanity_checker_func)
+    add_credentials_arguments(sanity_check_parser)
 
-    # Generic arguments
-    parser.add_argument("username", help="NovelAI username")
-    parser.add_argument("password", help="NovelAI password")
+    # Decode
+    decode_parser = subparser.add_parser("decode", help="Decode")
+    decode_parser.add_argument("model", help="Model to use")
+    decode_parser.add_argument("data", help="Data to decode")
+    decode_parser.set_defaults(func=decode_func)
 
     # Parse arguments
     args = parser.parse_args()
     if getattr(args, "func", None) is None:
         parser.print_help()
-    elif asyncio.iscoroutinefunction(args.func):
-        asyncio.run(args.func(args.username, args.password))
+        sys.exit(1)
+
+    # Get the values of the arguments
+    arg_names = inspect.getfullargspec(args.func).args
+    args_values = {name: getattr(args, name) for name in arg_names}
+
+    if asyncio.iscoroutinefunction(args.func):
+        asyncio.run(args.func(**args_values))
     else:
-        args.func(args.username, args.password)
+        args.func(**args_values)
diff --git a/tools/decode.py b/tools/decode.py
deleted file mode 100644
index 2103e42..0000000
--- a/tools/decode.py
+++ /dev/null
@@ -1,17 +0,0 @@
-from base64 import b64decode
-from os.path import abspath, dirname, join
-from sys import argv, path
-
-# pylint: disable=C0413,C0415
-path.insert(0, abspath(join(dirname(__file__), "..")))
-from novelai_api.Preset import Model
-from novelai_api.Tokenizer import Tokenizer
-
-assert 2 <= len(argv), "Expected argument"
-
-tokens = b64decode(argv[1])
-tokens = [int.from_bytes(tokens[i * 2 : (i + 1) * 2], "little") for i in range(len(tokens) // 2)]
-print(f"Tokens = {tokens}")
-
-text = Tokenizer.decode(Model.HypeBot, tokens)
-print(f"Text = {text}")