diff --git a/config.py b/config.py index cd7db83..7e8fa5d 100644 --- a/config.py +++ b/config.py @@ -87,6 +87,10 @@ class Config: def export_type(self): return self.get_arg('type', 'epub,txt') + @cached_property + def force(self): + return getattr(self._args, 'force', False) + def get_export_book(self, book, ext): temp = self.export_book_template for k in book.keys(): diff --git a/db.py b/db.py index c45e497..04b0cb1 100644 --- a/db.py +++ b/db.py @@ -1,6 +1,6 @@ import sqlite3 from semver import Version -from typing import Optional, Set, List +from typing import Dict, Optional, List VERSION_TABLE = '''CREATE TABLE version ( @@ -67,9 +67,9 @@ class CwmDb: self._db.execute('INSERT OR REPLACE INTO key VALUES (?, ?, ?);', [ chapter_id, user_id, key]) - def get_all_keys_as_origin(self) -> Set[str]: - cur = self._db.execute('SELECT chapter_id, user_id FROM key;') - return {f'{i[0]}{i[1]}' for i in cur} + def get_all_keys_as_origin(self) -> Dict[str, str]: + cur = self._db.execute('SELECT chapter_id, user_id, key FROM key;') + return {f'{i[0]}{i[1]}': i[2] for i in cur} def get_key(self, chapter_id: int) -> List[str]: cur = self._db.execute('SELECT key FROM key WHERE chapter_id = ?;', [ diff --git a/export.py b/export.py index f0ba743..5c31856 100644 --- a/export.py +++ b/export.py @@ -13,15 +13,18 @@ from random import choice key_imported = False +key_force_imported = False def get_key(db: CwmDb, cfg: Config, chapter_id: int): + global key_imported keys = db.get_key(chapter_id) if len(keys) == 0: if key_imported: raise ValueError('The key is not found.') else: import_keys(cfg.key, db) + key_imported = True keys = db.get_key(chapter_id) if len(keys) == 0: raise ValueError('The key is not found.') @@ -29,12 +32,17 @@ def get_key(db: CwmDb, cfg: Config, chapter_id: int): def try_decrypt(db: CwmDb, cfg: Config, content, chapter_id: int): + global key_force_imported keys = get_key(db, cfg, chapter_id) for key in keys: try: return decrypt(content, key).decode() except Exception: pass + if not key_force_imported: + import_keys(cfg.key, db, True) + key_force_imported = True + return try_decrypt(db, cfg, content, chapter_id) raise ValueError('Failed to decrypt the content.') diff --git a/key.py b/key.py index 30687a7..781c9a6 100644 --- a/key.py +++ b/key.py @@ -5,7 +5,7 @@ from zipfile import ZipFile from base64 import b64decode -def import_keys(key: str, db: CwmDb): +def import_keys(key: str, db: CwmDb, force=False): is_zip = False file_list = [] contain_dir_name = False @@ -25,7 +25,8 @@ def import_keys(key: str, db: CwmDb): for i in file_list: oid = b64decode(i).decode() if oid in keys: - continue + if not force: + continue cid = int(oid[0:9]) uid = int(oid[9:]) if is_zip: @@ -36,6 +37,8 @@ def import_keys(key: str, db: CwmDb): content = z.read(path).decode() else: content = open(join(key, i), 'r', encoding='UTF-8').read() + if oid in keys and content == keys[oid]: + continue db.add_key(cid, uid, content) count += 1 print(f'Imported {count} keys.') diff --git a/main.py b/main.py index 5b42316..03eaba0 100644 --- a/main.py +++ b/main.py @@ -22,6 +22,7 @@ parser.add_argument('--icd', '--image-cache-dir', help='Path to image cache dire parser.add_argument('-s', '--page-size', help='Maximum size of a page when asking for choices.', type=int, metavar='SIZE') # noqa: E501 parser.add_argument('-a', '--export-nodownload', help='export not downloaded chapter when exporting book.', type=parse_bool, metavar='BOOL') # noqa: E501 parser.add_argument('-i', '--image-type', help='How to handle images in EPUB. Available types: inline, footnote. Default: inline', choices=['inline', 'footnote'], metavar='TYPE') # noqa: E501 +parser.add_argument('-f', '--force', help='Force import keys.', action='store_true') # noqa: E501 parser.add_argument('action', help='The action to do.', choices=['importkey', 'exportchapter', 'exportbook', 'export', 'exportall', 'ik', 'ec', 'eb', 'e', 'ea'], nargs='?', default='export') # noqa: E501 @@ -40,7 +41,7 @@ def main(args=None): if cfg.key is None: raise ValueError('The key is not specified.') from key import import_keys - import_keys(cfg.key, db) + import_keys(cfg.key, db, cfg.force) elif arg.action == 'exportchapter' or arg.action == 'ec': if cfg.cwmdb is None: raise ValueError('The cwmdb is not specified.')