better error handling, split some files

CelestiaSunrise.py

@@ -4,14 +4,13 @@
 from __future__ import print_function, absolute_import
 import os
 import sys
-import traceback
 from src import (SaveManager,
                  decompress_data, compress_data,
                  XmlHandler, PonyShell)
 from docopt import docopt
 
 PRGM = os.path.basename(__file__)
-VERSION = "0.4.1b"
+VERSION = "0.4.2a"
 
 __doc__ = """
 {prgm} {ver}
@@ -41,9 +40,9 @@
     try:
         PonyShell(opts['<save_file>'], opts['<encrypt_key>']).cmdloop(intro=__doc__)
     except Exception as e:
-        print('Something went wrong, traceback:\n',
+        print('Something went wrong, error message:',
+              file=sys.stderr)
+        print('{}("{}")'.format(e.__class__.__name__, str(e)),
               file=sys.stderr)
-        tb = traceback.format_exc()
-        print(tb, file=sys.stderr)
     print('Exiting...')
     sys.exit(0)

src/ponyshell.py

@@ -8,131 +8,31 @@
 
 import base64
 import sys
+import binascii
 from cmd import Cmd
 from pprint import pprint
 from src import SaveManager, SaveError
 from src import decompress_data, compress_data
 from src import XmlHandler
 from src.docopt_utils import docopt_cmd
 from src.utility import Pony
+from src.show import *
+from src.set import *
 
-def show_currencies(xml_handle, args):
-    print('Main currencies:')
-    for cur, val in xml_handle.currencies['Main'].items():
-        print('  {}: {}'.format(cur, val))
-    print('\nShards:')
-    for cur, val in xml_handle.currencies['Shards'].items():
-        print('  {} shards: {}'.format(cur, val))
-    print('\nZecora ingredients:')
-    for cur, val in xml_handle.currencies['Ingredients'].items():
-        print('  {}: {}'.format(cur, val))
+class PonyError(Exception):
+    pass
 
-def show_currency(xml_handle, args):
-    for currency_id in args['<currency_id>']:
-        for typ in xml_handle.currencies.values():
-            for val in typ.values():
-                if currency_id == val.name:
-                    print(val)
-                    return
-
-def show_ponies(xml_handle, args):
-    print('Ponies:')
-    for pony in xml_handle.ponies.values():
-        print('  {}'.format(pony))
-    if args['-i']:
-        print('\nInventory ponies:')
-        for pony in xml_handle.inventory_ponies.values():
-            print('  {}'.format(pony))
-
-def show_pony(xml_handle, args):
-    for pony_id in args['<pony_id>']:
-        if pony_id in xml_handle.ponies:
-            print(xml_handle.ponies[pony_id])
-        if pony_id in xml_handle.inventory_ponies:
-            print(xml_handle.inventory_ponies[pony_id])
-
-def show_zones(xml_handle, args):
-    print('Zones:')
-    for zone in xml_handle.zones.values():
-        print('  {}'.format(zone))
-
-def show_zone(xml_handle, args):
-    for zone_id in args['<zone_id>']:
-        if zone_id in xml_handle.zones:
-            print(xml_handle.zones[zone_id])
-
-def set_currency(xml_handle, args):
-    for currency_id in args['<currency_id>']:
-        for typ in xml_handle.currencies.values():
-            for val in typ.values():
-                if currency_id == val.name:
-                    try:
-                        val.value = args['<value>']
-                    except ValueError as e:
-                        print(str(e))
-
-def process_set_pony(xml_handle, pony, args):
-    if args['level']:
-        if args['up']:
-            pony.level_up()
-        elif args['down']:
-            pony.level_down()
-        else:
-            pony.level = args['<value>']
-    elif args['shards']:
-        if args['up']:
-            pony.shard_up()
-        elif args['down']:
-            pony.shard_down()
-        else:
-            pony.shards = args['<value>']
-    elif args['next_game']:
-        try:
-            pony.next_game = Pony.GameTypes.map[args['<value>']]
-        except KeyError:
-            raise ValueError("Invalid game type")
-    elif args['reset_game_timer']:
-        pony.reset_game_timer()
-
-def set_ponies(xml_handle, args):
-    try:
-        for pony in xml_handle.ponies.values():
-            process_set_pony(xml_handle, pony, args)
-    except Exception as e:
-        print(str(e))
-
-def set_pony(xml_handle, args):
-    try:
-        for pony_id in args['<pony_id>']:
-            if pony_id in xml_handle.ponies:
-                process_set_pony(xml_handle, xml_handle.ponies[pony_id], args)
-    except Exception as e:
-        print(str(e))
-
-def process_set_zone(xml_handle, zone, args):
-    if args['clearables']:
-        zone.clear_clearable_items()
-    elif args['foes']:
-        zone.clear_foes()
-    else:
-        zone.clear_all()
-
-def set_zones(xml_handle, args):
-    for zone in xml_handle.zones.values():
-        process_set_zone(xml_handle, zone, args)
-
-def set_zone(xml_handle, args):
-    for zone_id in args['<zone_id>']:
-        if zone_id in xml_handle.zones:
-            process_set_zone(xml_handle, xml_handle.zones[zone_id], args)
 
 class PonyShell(Cmd):
 
     prompt = 'ponyshell> '
 
     def __init__(self, savefile, gluid):
         Cmd.__init__(self)
-        self._save_manager = SaveManager(savefile, base64.b64decode(gluid))
+        try:
+            self._save_manager = SaveManager(savefile, base64.b64decode(gluid))
+        except binascii.Error:
+            raise PonyError("Invalid decryption key")
         self._xml_handle = XmlHandler(decompress_data(self._save_manager.load())
                                       .decode('utf-8'))
         self._xml_handle.pre_load()
@@ -274,8 +174,11 @@ def do_write_save(self, args):
 Options:
   -h --help     Show this help."""
         if args['<gluid>'] is not None:
-            args['<gluid>'] = base64.b64decode(args['<gluid>'].encode('utf-8'))
-            print(args['<gluid>'])
+            try:
+                args['<gluid>'] = base64.b64decode(args['<gluid>'].encode('utf-8'))
+            except binascii.Error:
+                print("Invalid encryption key")
+                return
         try:
             self._save_manager.save(compress_data(repr(self._xml_handle)
                                                   .encode('utf-8')),
@@ -296,4 +199,4 @@ def emptyline(self):
         pass
 
     def default(self, line):
-        print("{}: command not found".format(line), file=sys.stderr)
+        print("{}: command not found".format(line))

src/savemanager.py

@@ -11,16 +11,16 @@
 import struct
 import src.xxtea as xxtea
 
+class SaveError(Exception):
+    pass
+
+
 def read_or_raise(file, size):
     data = file.read(size)
     if len(data) != size:
         raise SaveError("Unable to read, truncated or corrupted file")
     return data
 
-class SaveError(Exception):
-    pass
-
-
 class SaveManager(object):
     def __init__(self, filename, gluid):
         self.filename = filename
@@ -52,9 +52,8 @@ def _load_buffer(self, file):
         try:
             res = zlib.decompress(decrypt_data)
         except zlib.error as e:
-            raise SaveError(str(e))
-        else:
-            self.results.append(res)
+            raise SaveError("Unable to decompress data, truncated or corrupted file, or bad decryption key")
+        self.results.append(res)
         if len(res) != uncompress_size:
             raise SaveError("Invalid inflated data")
         crc_res = zlib.crc32(res)
@@ -96,10 +95,7 @@ def save(self, data, filename, gluid=None):
 
 
 def decompress_data(data):
-    try:
-        res = zlib.decompress(data[16:])
-    except zlib.error as e:
-        raise SaveError(str(e))
+    res = zlib.decompress(data[16:])
     uncompress_size, compress_size = struct.unpack('2I', data[:8])
     if len(res) != uncompress_size or len(data[16:]) != compress_size:
         raise SaveError("Invalid inflated data")

src/set.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# File: set.py
+# by Arzaroth Lekva
+# [email protected]
+#
+
+def set_currency(xml_handle, args):
+    for currency_id in args['<currency_id>']:
+        for typ in xml_handle.currencies.values():
+            for val in typ.values():
+                if currency_id == val.name:
+                    try:
+                        val.value = args['<value>']
+                    except ValueError as e:
+                        print(str(e))
+
+def _process_set_pony(xml_handle, pony, args):
+    if args['level']:
+        if args['up']:
+            pony.level_up()
+        elif args['down']:
+            pony.level_down()
+        else:
+            pony.level = args['<value>']
+    elif args['shards']:
+        if args['up']:
+            pony.shard_up()
+        elif args['down']:
+            pony.shard_down()
+        else:
+            pony.shards = args['<value>']
+    elif args['next_game']:
+        try:
+            pony.next_game = Pony.GameTypes.map[args['<value>']]
+        except KeyError:
+            raise ValueError("Invalid game type")
+    elif args['reset_game_timer']:
+        pony.reset_game_timer()
+
+def set_ponies(xml_handle, args):
+    try:
+        for pony in xml_handle.ponies.values():
+            _process_set_pony(xml_handle, pony, args)
+    except Exception as e:
+        print(str(e))
+
+def set_pony(xml_handle, args):
+    try:
+        for pony_id in args['<pony_id>']:
+            if pony_id in xml_handle.ponies:
+                _process_set_pony(xml_handle, xml_handle.ponies[pony_id], args)
+    except Exception as e:
+        print(str(e))
+
+def _process_set_zone(xml_handle, zone, args):
+    if args['clearables']:
+        zone.clear_clearable_items()
+    elif args['foes']:
+        zone.clear_foes()
+    else:
+        zone.clear_all()
+
+def set_zones(xml_handle, args):
+    for zone in xml_handle.zones.values():
+        _process_set_zone(xml_handle, zone, args)
+
+def set_zone(xml_handle, args):
+    for zone_id in args['<zone_id>']:
+        if zone_id in xml_handle.zones:
+            _process_set_zone(xml_handle, xml_handle.zones[zone_id], args)

src/show.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# File: show.py
+# by Arzaroth Lekva
+# [email protected]
+#
+
+def show_currencies(xml_handle, args):
+    print('Main currencies:')
+    for cur, val in xml_handle.currencies['Main'].items():
+        print('  {}: {}'.format(cur, val))
+    print('\nShards:')
+    for cur, val in xml_handle.currencies['Shards'].items():
+        print('  {} shards: {}'.format(cur, val))
+    print('\nZecora ingredients:')
+    for cur, val in xml_handle.currencies['Ingredients'].items():
+        print('  {}: {}'.format(cur, val))
+
+def show_currency(xml_handle, args):
+    for currency_id in args['<currency_id>']:
+        for typ in xml_handle.currencies.values():
+            for val in typ.values():
+                if currency_id == val.name:
+                    print(val)
+                    return
+
+def show_ponies(xml_handle, args):
+    print('Ponies:')
+    for pony in xml_handle.ponies.values():
+        print('  {}'.format(pony))
+    if args['-i']:
+        print('\nInventory ponies:')
+        for pony in xml_handle.inventory_ponies.values():
+            print('  {}'.format(pony))
+
+def show_pony(xml_handle, args):
+    for pony_id in args['<pony_id>']:
+        if pony_id in xml_handle.ponies:
+            print(xml_handle.ponies[pony_id])
+        if pony_id in xml_handle.inventory_ponies:
+            print(xml_handle.inventory_ponies[pony_id])
+
+def show_zones(xml_handle, args):
+    print('Zones:')
+    for zone in xml_handle.zones.values():
+        print('  {}'.format(zone))
+
+def show_zone(xml_handle, args):
+    for zone_id in args['<zone_id>']:
+        if zone_id in xml_handle.zones:
+            print(xml_handle.zones[zone_id])