Use key names for identification, add options to hash filenames

ttsmirror.py

@@ -1,14 +1,15 @@
 import os
 import logging
 import requests
-from urllib.parse import urlparse, urljoin
+from urllib.parse import urljoin
 import json
+import hashlib
 
 __version__ = '0.1'
 __author__ = 'Andrew Williams'
 
 
-def iterate_save(obj, output_path, url_prefix):
+def iterate_save(obj, output_path, url_prefix, hash_filename=False):
     """
     Iterate a save, download assets, and update the locations as needed.
 
@@ -25,15 +26,17 @@ def iterate_save(obj, output_path, url_prefix):
 
     for key, val in iter:
         if isinstance(val, dict) or isinstance(val, list):
-            res = iterate_save(val, output_path, url_prefix)
+            res = iterate_save(val, output_path, url_prefix, hash_filename)
             obj[key] = res
         if isinstance(val, str):
-            parsed = urlparse(val)
-            if parsed.scheme and parsed.scheme in ['http', 'https', 'ftp']:
+            if key.lower()[-3:] == 'url' and val.strip() != '':
                 # Generate new filename
-                new_filename = val
-                for rep in [':', '/', '?', '=']:
-                    new_filename = new_filename.replace(rep, '_')
+                if hash_filename:
+                    new_filename = hashlib.sha1(val.encode('utf-8')).hexdigest()
+                else:
+                    new_filename = val
+                    for rep in [':', '/', '?', '=']:
+                        new_filename = new_filename.replace(rep, '_')
                 # Check if exists
                 if not os.path.exists(os.path.join(output_path, new_filename)):
                     res = requests.get(val, stream=True)
@@ -48,18 +51,18 @@ def iterate_save(obj, output_path, url_prefix):
     return obj
 
 
-def process_save(filename, output_path, url_prefix):
+def process_save(filename, output_path, url_prefix, hash_filename):
     """Parses TTS JSON save file and mirrors the required objects."""
     new_filename = '%s_new.json' % filename.replace('.', '_')
     with open(filename, 'r') as fobj:
         save = json.load(fobj)
     try:
-        new_save = iterate_save(save, output_path, url_prefix)
+        new_save = iterate_save(save, output_path, url_prefix, hash_filename)
     except Exception as e:
         logging.exception('Unable to process save: %s' % e)
     new_save['SaveName'] = '%s - Mirrored' % new_save['SaveName']
     with open(new_filename, 'w') as outfobj:
-        outfobj.write(json.dumps(new_save))
+        outfobj.write(json.dumps(new_save, sort_keys=True, indent=4, separators=(',', ': ')))
     logging.debug('Done, wrote %s', new_filename)
 
 
@@ -72,6 +75,7 @@ def main():
     parser.add_argument('save_file')
     parser.add_argument('output_path', default=os.path.curdir)
     parser.add_argument('url_prefix')
+    parser.add_argument('--hash_filename', action='store_true', dest='hash_filename')
 
     logging.basicConfig(level=logging.DEBUG)
     args = parser.parse_args()
@@ -80,7 +84,7 @@ def main():
         logging.error('Unknown file %s' % args.save_file)
         sys.exit(1)
 
-    process_save(args.save_file, args.output_path, args.url_prefix)
+    process_save(args.save_file, args.output_path, args.url_prefix, args.hash_filename)
 
 
 if __name__ == '__main__':