diff --git a/CelestiaSunrise.py b/CelestiaSunrise.py index f915184..f1c793d 100755 --- a/CelestiaSunrise.py +++ b/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[''], opts['']).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) diff --git a/src/defaultordereddict.py b/src/defaultordereddict.py old mode 100755 new mode 100644 diff --git a/src/enum.py b/src/enum.py old mode 100755 new mode 100644 diff --git a/src/ponyshell.py b/src/ponyshell.py index 0095ed6..669259a 100644 --- a/src/ponyshell.py +++ b/src/ponyshell.py @@ -8,6 +8,7 @@ import base64 import sys +import binascii from cmd import Cmd from pprint import pprint from src import SaveManager, SaveError @@ -15,116 +16,12 @@ 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['']: - 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['']: - 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['']: - if zone_id in xml_handle.zones: - print(xml_handle.zones[zone_id]) - -def set_currency(xml_handle, args): - for currency_id in args['']: - for typ in xml_handle.currencies.values(): - for val in typ.values(): - if currency_id == val.name: - try: - val.value = args[''] - 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[''] - elif args['shards']: - if args['up']: - pony.shard_up() - elif args['down']: - pony.shard_down() - else: - pony.shards = args[''] - elif args['next_game']: - try: - pony.next_game = Pony.GameTypes.map[args['']] - 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['']: - 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['']: - if zone_id in xml_handle.zones: - process_set_zone(xml_handle, xml_handle.zones[zone_id], args) class PonyShell(Cmd): @@ -132,7 +29,10 @@ class PonyShell(Cmd): 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[''] is not None: - args[''] = base64.b64decode(args[''].encode('utf-8')) - print(args['']) + try: + args[''] = base64.b64decode(args[''].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)) diff --git a/src/savemanager.py b/src/savemanager.py index 2049f80..3be1cf4 100644 --- a/src/savemanager.py +++ b/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") diff --git a/src/set.py b/src/set.py new file mode 100644 index 0000000..c49bb1c --- /dev/null +++ b/src/set.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# File: set.py +# by Arzaroth Lekva +# arzaroth@arzaroth.com +# + +def set_currency(xml_handle, args): + for currency_id in args['']: + for typ in xml_handle.currencies.values(): + for val in typ.values(): + if currency_id == val.name: + try: + val.value = args[''] + 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[''] + elif args['shards']: + if args['up']: + pony.shard_up() + elif args['down']: + pony.shard_down() + else: + pony.shards = args[''] + elif args['next_game']: + try: + pony.next_game = Pony.GameTypes.map[args['']] + 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['']: + 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['']: + if zone_id in xml_handle.zones: + _process_set_zone(xml_handle, xml_handle.zones[zone_id], args) diff --git a/src/show.py b/src/show.py new file mode 100644 index 0000000..cb35083 --- /dev/null +++ b/src/show.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# File: show.py +# by Arzaroth Lekva +# arzaroth@arzaroth.com +# + +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['']: + 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['']: + 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['']: + if zone_id in xml_handle.zones: + print(xml_handle.zones[zone_id])