add new data, improve script

1 files changed, 41 insertions, 18 deletions

diff --git a/main.py b/main.py
index ed0d793..0d3c6d5 100755
--- a/main.py
+++ b/main.py
@@ -153,12 +153,31 @@ def decode2(s):
         buf += letter
 
     return buf
+
+
+# s: source
+# t: type
+def decode_auto(s, t, reverse_decoded=False, remove_junk=True):
+    if t == 1:
+        s = clean_string(s, remove_junk=remove_junk)
+        result = decode(s)
+    
+    elif t == 2:
+        result = decode2(s)
+
+    if reverse_decoded:
+        # reverse string
+        result = result[::-1]
+
+    return result
+
         
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--decode', action='store_true')
     parser.add_argument('--stats', action='store_true')
     parser.add_argument('--decode-string')
+    parser.add_argument('--decode-file')
     parser.add_argument('--with-junk', action='store_true')
     parser.add_argument('--is-url', action='store_true')
     parser.add_argument('--type', type=int, choices=[1, 2], default=1)
@@ -179,18 +198,12 @@ def main():
         
         for obj in data:
             text = obj['text']
-            
-            if args.type == 1:
-                text = clean_string(text, remove_junk=(not args.with_junk))
-                text_decoded = decode(text)
-            
-            elif args.type == 2:
-                text_decoded = decode2(text)
-
-            if args.reverse_decoded:
-                # reverse string
-                text_decoded = text_decoded[::-1]
-
+            text_decoded = decode_auto(text,
+                                      args.type,
+                                      remove_junk=(not args.with_junk),
+                                      reverse_decoded=args.reverse_decoded)
+         
+            # print all information
             print(obj['text'])
             print_colored(text, 'green', fallback_prefix='[CLEANED] ')
             print_colored(text_decoded, 'cyan', fallback_prefix='[DECODED] ')
@@ -198,15 +211,25 @@ def main():
             if 'pic' in obj:
                 pic = obj['pic'] if isinstance(obj['pic'], list) else [obj['pic']]
                 print_colored(', '.join(pic), 'red', fallback_prefix='[PICS] ')
+            if 'link' in obj:
+                print_colored(obj['link'], 'red', fallback_prefix='[LINK] ')
 
             print("\n")
 
-    elif args.decode_string:
-        text = clean_string(args.decode_string, remove_junk=(not args.with_junk))
-
-        print(args.decode_string)
-        print_colored(text, 'green', fallback_prefix='[CLEANED] ')
-        print_colored(decode(text, is_url=args.is_url), 'cyan', fallback_prefix='[DECODED] ')
+    elif args.decode_string or args.decode_file:
+        if args.decode_string:
+            source = args.decode_string
+        else: 
+            with open(args.decode_file, 'r') as f:
+                source = f.read()
+
+        text_decoded = decode_auto(source,
+                                   args.type,
+                                   remove_junk=(not args.with_junk),
+                                   reverse_decoded=args.reverse_decoded)
+        
+        # print
+        print_colored(text_decoded, 'cyan', fallback_prefix='[DECODED] ')
             
     elif args.stats:
         count = len(data)