Implement decrypting of all books

1 files changed, 15 insertions, 11 deletions

diff --git a/Other_Tools/Kobo/obok.py b/Other_Tools/Kobo/obok.py
index 7a50ba9..37ee286 100644
--- a/Other_Tools/Kobo/obok.py
+++ b/Other_Tools/Kobo/obok.py
@@ -674,7 +674,7 @@ def decrypt_book(book, lib):
         print u"DRM-free book, conversion is not needed"
         shutil.copyfile(book.filename, outname)
         print u"Book saved as {0}".format(os.path.join(os.getcwd(), outname))
-        exit(0)
+        return 0
     result = 1
     for userkey in lib.userkeys:
         print u"Trying key: {0}".format(userkey.encode('hex_codec'))
@@ -718,19 +718,23 @@ def cli_main():
         print u"{0}: {1}".format(i + 1, book.title)
     print u"Or 'all'"
 
-    num_string = raw_input(u"Convert book number... ")
-    try:
-        num = int(num_string)
-        book = lib.books[num - 1]
-    except (ValueError, IndexError):
-        print u"Invalid choice. Exiting..."
-        exit()
+    choice = raw_input(u"Convert book number... ")
+    if choice == u'all':
+        books = list(lib.books)
+    else:
+        try:
+            num = int(choice)
+            books = [lib.books[num - 1]]
+        except (ValueError, IndexError):
+            print u"Invalid choice. Exiting..."
+            exit()
 
-    result = decrypt_book(book, lib)
+    results = [decrypt_book(book, lib) for book in books]
     lib.close()
-    if result != 0:
+    overall_result = all(result != 0 for result in results)
+    if overall_result != 0:
         print u"Could not decrypt book with any of the keys found."
-    return result
+    return overall_result
 
 
 if __name__ == '__main__':