diff --git a/example/boilerplate.py b/example/boilerplate.py index 538c3fd..8156769 100644 --- a/example/boilerplate.py +++ b/example/boilerplate.py @@ -1,10 +1,12 @@ import json +from datetime import datetime from logging import Logger, StreamHandler from os import environ as env from pathlib import Path from typing import Any, Optional from aiohttp import ClientSession +from msgpackr.constants import UNDEFINED from novelai_api import NovelAIAPI from novelai_api.utils import get_encryption_key @@ -85,6 +87,10 @@ class JSONEncoder(json.JSONEncoder): def default(self, o: Any) -> Any: if isinstance(o, bytes): return o.hex() + if o is UNDEFINED: + return "" + if isinstance(o, datetime): + return o.isoformat() return super().default(o) diff --git a/example/download_last_story_and_decrypt.py b/example/download_last_story_and_decrypt.py new file mode 100644 index 0000000..5cd5f4e --- /dev/null +++ b/example/download_last_story_and_decrypt.py @@ -0,0 +1,30 @@ +import asyncio +from pathlib import Path + +from example.boilerplate import API, dumps +from novelai_api.utils import decrypt_user_data + +dump_file = Path(__file__).parent.parent / "results" / "story.json" + + +async def main(): + async with API() as api_handler: + api = api_handler.api + key = api_handler.encryption_key + + keystore = await api.high_level.get_keystore(key) + + story = (await api.high_level.download_user_stories())[0] + decrypt_user_data(story, keystore) + + storycontent_id = story["data"]["remoteStoryId"] + + story_contents = await api.low_level.download_object("storycontent", storycontent_id) + decrypt_user_data(story_contents, keystore, True) + + dump_file.parent.mkdir(exist_ok=True) + dump_file.write_text(dumps(story_contents), "utf-8") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/novelai_api/GlobalSettings.py b/novelai_api/GlobalSettings.py index 0fd2bd3..9889041 100644 --- a/novelai_api/GlobalSettings.py +++ b/novelai_api/GlobalSettings.py @@ -712,14 +712,6 @@ class GlobalSettings: ".", ":", "\n", - "ve", - "s", - "t", - "n", - "d", - "ll", - "re", - "m", "-", "*", ")", diff --git a/novelai_api/Msgpackr_Extensions.py b/novelai_api/Msgpackr_Extensions.py new file mode 100644 index 0000000..d99cfc8 --- /dev/null +++ b/novelai_api/Msgpackr_Extensions.py @@ -0,0 +1,40 @@ +from typing import Any + +from msgpackr.constants import SKIP +from msgpackr.unpack import MsgpackExtension + + +class NAIExtension(MsgpackExtension): + @classmethod + def unpack(cls, _unpacker, _data: bytes, _pos: int, _length: int) -> Any: + return SKIP + + # TODO: the data should be bundled in post_unpack + + @classmethod + def pack(cls, _unpacker, data: Any) -> bytes: + return data + + +class Ext20(NAIExtension): + EXT_TYPE = 20 + + +class Ext30(NAIExtension): + EXT_TYPE = 30 + + +class Ext31(NAIExtension): + EXT_TYPE = 31 + + +class Ext40(NAIExtension): + EXT_TYPE = 40 + + +class Ext41(NAIExtension): + EXT_TYPE = 41 + + +class Ext42(NAIExtension): + EXT_TYPE = 42 diff --git a/novelai_api/utils.py b/novelai_api/utils.py index 17fead3..3e09ec4 100644 --- a/novelai_api/utils.py +++ b/novelai_api/utils.py @@ -7,15 +7,21 @@ from zlib import decompress as inflate import argon2 +from msgpackr import Unpacker from nacl.exceptions import CryptoError from nacl.secret import SecretBox from novelai_api.Keystore import Keystore +from novelai_api.Msgpackr_Extensions import Ext20, Ext30, Ext31, Ext40, Ext41, Ext42 from novelai_api.NovelAIError import NovelAIError from novelai_api.Preset import Model, Preset from novelai_api.python_utils import assert_type from novelai_api.Tokenizer import Tokenizer +unpacker = Unpacker() +unpacker.register_extensions(Ext20, Ext30, Ext31, Ext40, Ext41, Ext42) +unpacker_state = unpacker.export_state() + # API utils def argon_hash(email: str, password: str, size: int, domain: str) -> str: @@ -185,13 +191,16 @@ def compress_user_data(items: Union[List[Dict[str, Any]], Dict[str, Any]]): del item["decrypted"] -def decrypt_user_data(items: Union[List[Dict[str, Any]], Dict[str, Any]], keystore: Keystore): +def decrypt_user_data( + items: Union[List[Dict[str, Any]], Dict[str, Any]], keystore: Keystore, uncompress_document: bool = False +): """ Decrypt the data of each item in :ref: items If an item has already been decrypted, it won't be decrypted a second time :param items: Item or list of items to decrypt :param keystore: Keystore retrieved with the get_keystore method + :param uncompress_document: If True, the document will be decompressed """ # 1 item @@ -224,6 +233,11 @@ def decrypt_user_data(items: Union[List[Dict[str, Any]], Dict[str, Any]], keysto item["nonce"] = nonce item["decrypted"] = True item["compressed"] = is_compressed + + if uncompress_document and "document" in data: + unpacker.restore_state(unpacker_state) + data["document"] = unpacker.unpack(b64decode(data["document"])) + continue except json.JSONDecodeError: diff --git a/pyproject.toml b/pyproject.toml index fd84900..b633af1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ ftfy = "^6.1.3" regex = "^2023.12.25" sentencepiece = "^0.2.0" poetry = "^1.7.1" -# msgpack = "^1.0.5" +msgpackr-python = "^0.1.2" [tool.poetry.group.dev.dependencies] python-dotenv = "^0.21.1"