Add support to force import keys

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():

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 = ?;', [

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.')
 
 

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.')

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.')