class SaveManager(object):
    SECTIONS = 2

    def __init__(self, filename, gluid=b''):
        self.filename = filename
        self.gluid = gluid

    def _save_buffer(self, data, file, gluid=None):
        crc_res = zlib.crc32(data) & 0xffffffff
        uncompress_size = len(data)
        decrypt_data = zlib.compress(data)
        crc_offset = len(decrypt_data) + 4
        decrypt_data += struct.pack('I', crc_res)
        raw_data = xxtea.encrypt(decrypt_data,
                                 self.gluid if gluid is None else gluid)
        data_size = len(raw_data)
        file.write(struct.pack('3I', uncompress_size, crc_offset, data_size))
        file.write(raw_data)

    def _load_buffer(self, file):
        metadata = read_or_raise(file, 12)
        uncompress_size, crc_offset, data_size = struct.unpack('3I', metadata)
        if crc_offset > data_size:
            raise SaveError("Bad size or crc_offset")
        raw_data = read_or_raise(file, data_size)
        decrypt_data = xxtea.decrypt(raw_data,
                                     self.gluid)
        crc_value = struct.unpack('I', decrypt_data[crc_offset - 4:crc_offset])[0]
        try:
            res = zlib.decompress(decrypt_data)
        except zlib.error as e:
            raise SaveError("Unable to decompress data, truncated or corrupted file, or bad decryption key")
        if len(res) != uncompress_size:
            raise SaveError("Invalid inflated data")
        crc_res = zlib.crc32(res) & 0xffffffff
        if crc_res != crc_value:
            raise SaveError("crc mismatch")
        return res

    def load(self, legacy=False):
        try:
            with open(self.filename, 'rb') as file:
                print('Loading save.', end='')
                if legacy:
                    results = [decompress_data(file.read())]
                    print('.', end='')
                else:
                    file.seek(-4, 2)
                    if struct.unpack('I', read_or_raise(file, 4))[0] != SaveManager.SECTIONS:
                        raise SaveError("Invalid sections number, truncated or corrupted file")
                    print('.', end='')
                    file.seek(0, 0)
                    results = [self._load_buffer(file) for _ in range(SaveManager.SECTIONS)]
                print('.\nDone!')
        except Exception as e:
            raise SaveError(str(e))
        save_number = struct.unpack('I', results[0])[0] if len(results) > 1 else 10
        return results[-1], save_number

    def save(self, data, filename, save_number=10,
             gluid=None, legacy=False):
        try:
            with open(filename, 'wb') as file:
                print('Writing save.', end='')
                if legacy:
                    file.write(data)
                    print('.', end='')
                else:
                    self._save_buffer(struct.pack('I', save_number), file, gluid)
                    print('.', end='')
                    self._save_buffer(data, file, gluid)
                    print('.', end='')
                    file.write(struct.pack('I', SaveManager.SECTIONS))
                print('.\nDone!')
        except Exception as e:
            raise SaveError(str(e)) class SaveError(Exception):
    pass def decompress_data(data):
    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")
    return res def compress_data(data):
    compressed_data = zlib.compress(data)
    res = struct.pack('4I', len(data), len(compressed_data), 0, 0)
    res += compressed_data
    return res class XmlHandler(object):
    ponies = XmlDescriptor()
    inventory = XmlDescriptor()
    missing_ponies = XmlDescriptor()
    currencies = XmlDescriptor()
    player_infos = XmlDescriptor()
    actions = XmlDescriptor()
    zones = XmlDescriptor()
    quests = XmlDescriptor()
    _mapzones = XmlDescriptor()

    def __init__(self, xml_data):
        print('Parsing XML tree...')
        self.xmlobj = rapidxml.RapidXml(bytearray(xml_data))

    def _get__mapzones(self):
        if type(self.xmlobj['MLP_Save']['MapZone']) != list:
            return [self.xmlobj['MLP_Save']['MapZone']]
        return self.xmlobj['MLP_Save']['MapZone']

    def _filtered_actions(self, ID):
        res = defaultdict(dict)
        for typ, actions in self.actions['Ponies'].items():
            for action, tags in actions.items():
                items = tags[0]['Item']
                if type(items) != list:
                    items = [items]
                tag = [i for i in items if i['@Category'].value == ID]
                if not tag:
                    tag = self.xmlobj.allocate_node('Item')
                    tag.append_attribute(self.xmlobj.allocate_attribute('Category', ID))
                    tag.append_attribute(self.xmlobj.allocate_attribute('Value', '0'))
                    tags[0].append_node(tag)
                    tag = [tag]
                res[typ][action] = tag
        return {'Pony': res, 'Global': self.actions['Global']}

    def _get_ponies(self):
        res = OrderedDict()
        for mapzone in self._mapzones:
            try:
                ponyobjects = mapzone['GameObjects']['Pony_Objects']
            except KeyError:
                pass
            else:
                for ponytag in ponyobjects:
                    ID = ponytag["@ID"].value
                    res[ID] = Pony(ponytag, self._filtered_actions(ID),
                                   None if ID not in PONY_LIST else
                                   PONY_LIST[ID])
        return res

    def _get_inventory(self):
        try:
            storage = self.xmlobj['MLP_Save']['PlayerData']['Storage']
        except KeyError:
            storage = self.xmlobj.allocate_node('Storage')
            self.xmlobj['MLP_Save']['PlayerData'].append_node(storage)
        return Inventory(storage, self.xmlobj)

    def _get_missing_ponies(self):
        return MissingPonies(self.ponies, self.inventory.ponies, PONY_LIST)

    def _get_currencies(self):
        playerdata = self.xmlobj['MLP_Save']['PlayerData']
        res = DefaultOrderedDict(OrderedDict)
        main = res['Main currencies']
        main['Bits'] = Currency('@Coins', 'Bits', playerdata)
        main['Gems'] = Currency('@Hearts', 'Gems', playerdata)
        main['Hearts'] = Currency('@Social', 'Hearts', playerdata)
        shards = playerdata['Shards']
        for i in ('Loyalty', 'Honesty', 'Kindness', 'Generosity', 'Laughter', 'Magic'):
            res['Shards'][i + ' shards'] = Currency('@' + i, i + ' shards', shards, 999)
        try:
            # Minecart Update (1.8)
            main['Wheels'] = Currency('@Wheels', 'Wheels', playerdata['Minecart'], 5)
        except:
            pass
        try:
            # Dragon Update (2.3)
            main['Sapphires'] = Currency('@BossEventCurrency', 'Sapphires', playerdata)
        except:
            pass
        try:
            # Everfree Update (2.1)
            ingredients = playerdata['Ingredients']
            zecora = res['Zecora ingredients']
            zecora['Black Iris'] = Currency('@BlackIris', 'Black Iris', ingredients, 5)
            zecora['Garlic'] = Currency('@Garlic', 'Garlic', ingredients, 5)
            zecora['Sticky Sap'] = Currency('@GlueTree', 'Sticky Sap', ingredients, 5)
            zecora['Joke Plant'] = Currency('@PoisonJokePlant', 'Joke Plant', ingredients, 5)
            zecora['Purple Mushrooms'] = Currency('@PurpleGlowingMushrooms', 'Purple Mushrooms',
                                                  ingredients, 5)
            zecora['Red Orchid'] = Currency('@RedOrchid', 'Red Orchid', ingredients, 5)
        except:
            pass
        try:
            # Boutique Update (2.6)
            popcurrency = playerdata['PopCurrency']
            boutique = res['Boutique ingredients']
            boutique['Flower'] = Currency('@PopCurrency1', 'Flower', popcurrency, 999)
            boutique['Button'] = Currency('@PopCurrency2', 'Button', popcurrency, 999)
            boutique['Thread'] = Currency('@PopCurrency3', 'Thread', popcurrency, 999)
            boutique['Fabric'] = Currency('@PopCurrency4', 'Fabric', popcurrency, 999)
            boutique['Ribbon'] = Currency('@PopCurrency5', 'Ribbon', popcurrency, 999)
        except:
            pass
        try:
            # Party Update (2.8)
            main['Party Points'] = Currency('@SpecialCurrency', 'Party Points', playerdata)
        except:
            pass
        try:
            # Countess Coloratura Update (3.0)
            tasktokens = playerdata['TaskTokens']
            tokens = res['Tokens']
            tokens['Camera'] = Currency('@Token_Camera', 'Camera', tasktokens)
            tokens['Leather Jacket'] = Currency('@Token_Leather_Jacket', 'Leather Jacket', tasktokens)
            tokens['Hair Dye'] = Currency('@Token_Hair_dye', 'Hair Dye', tasktokens)
            tokens['Note'] = Currency('@Token_Note', 'Note', tasktokens)
            tokens['Glasses'] = Currency('@Token_Glasses', 'Glasses', tasktokens)
            tokens['Microphone'] = Currency('@Token_Microphone', 'Microphone', tasktokens)
            tokens['Cap'] = Currency('@Token_Cap', 'Cap', tasktokens)
        except:
            pass
        return res

    def _get_player_infos(self):
        playerdata = self.xmlobj['MLP_Save']['PlayerData']
        res = OrderedDict()
        res['Level'] = PlayerData('@Level', 'Level', playerdata, 135)
        res['XP'] = PlayerData('@XP', 'XP', playerdata)
        try:
            # VIP Update (2.7)
            res['VIP Points'] = PlayerData('@vip_points', 'VIP Points', playerdata['vip'])
        except:
            pass
        return res

    def _get_zones(self):
        mapzones_spec = OrderedDict((
            ("0", {"name": "Ponyville",
                   "foes": {"ID": "Parasprite_Objects",
                            "name": "Parasprites"}}),
            ("1", {"name": "Canterlot",
                   "foes": {"ID": "Changeling_Objects",
                            "name": "Changelings"}}),
            ("2", {"name": "Sweet Apple Acres",
                   "foes": {"ID": "Parasprite_Objects",
                            "name": "Vampire Bats"}}),
            ("3", {"name": "Everfree Forest",
                   "foes": {"ID": "Plunderseed_Vine_Objects",
                            "name": "Plunderseed Vines"}}),
        ))
        zones = OrderedDict()
        for mapzone in self._mapzones:
            gameobjects = mapzone['GameObjects']
            try:
                zone_spec = mapzones_spec[mapzone["@ID"].value]
            except KeyError:
                continue
            clearables = Clearables('Clearable_Objects',
                                    gameobjects)
            foes = Foes(zone_spec["foes"]["ID"],
                        zone_spec["foes"]["name"],
                        gameobjects)
            shops = Shops(gameobjects['Pony_House_Objects'])
            zones[mapzone["@ID"].value] = Zone(mapzone["@ID"].value,
                                               zone_spec["name"],
                                               clearables,
                                               foes, shops)
        return zones

    def _get_actions(self):
        datatable = self.xmlobj['MLP_Save']['QuestData']['GlobalDataTable']['DataTable']
        objectcategories = datatable['ObjectCategoryList']['ObjectCategory']
        globalcategories = datatable['GlobalCategoryList']['GlobalCategory']
        actions = defaultdict(lambda: defaultdict(list))
        glob = defaultdict(lambda: defaultdict(list))
        def populate_dict(dikt, key='@ID', suffix='', inner=None):
            for typ in ('Complete', 'Started', 'ItemSelected'):
                for action in Pony.GameTypes.map:
                    if (tag[key].value == ('PlayAction%s_%s%s' % (typ, action, suffix))
                        and (not inner or inner in tag)):
                        dikt[typ][action].append(tag)
            for typ in ('Complete', 'Started'):
                if tag[key].value == ('ClearSkies_%s%s' % (typ, suffix)):
                    dikt['ClearSkies'][typ].append(tag)
            if tag[key].value == ('PlayActionComplete%s' % suffix):
                dikt['Complete']['All'].append(tag)

        for tag in objectcategories:
            populate_dict(actions, inner='Item')
        for tag in globalcategories:
            for suffix in (' [TOTAL]', ' [TOTAL] Pony'):
                populate_dict(glob, key='@Category', suffix=suffix)
        return {'Global': glob, 'Ponies': actions}

    def _get_quests(self):
        activequestlist = self.xmlobj['MLP_Save']['QuestData']['ActiveQuestList']
        return [Quest(quest) for quest in activequestlist]

    def pre_load(self):
        self.player_infos
        self.currencies
        self.ponies
        self.inventory
        self.missing_ponies
        self.zones
        self.quests
        self.actions

    def to_string(self):
        return self.xmlobj.unparse(raw=True)

    def prettify(self):
        return self.xmlobj.unparse(pretty=True, raw=True) class Pony(StoredItem):
    GameTypes = enum(Ball=0, Apple=1, Book=2)

    def __init__(self, tag, actions, name=None):
        StoredItem.__init__(self, tag)
        if name is not None:
            self.name = name
        self._leveltag = tag["Game"]["Level"]
        self._minigametag = tag["Game"]["MiniGame"]
        self._pony_actions = actions
        self.stored = False

    @StoredItem._level.getter
    def _level(self):
        return min(max(int(self._leveltag["@Level"].value), 0), 5)

    @StoredItem._exp.getter
    def _exp(self):
        return int(self._leveltag["@CurrentEXP"].value)

    @StoredItem._shards.getter
    def _shards(self):
        return min(max(int(self._leveltag["@Shards"].value), 0), 10)

    @StoredItem.level.setter
    def level(self, new_val):
        new_val = int(new_val)
        if new_val < 0 or new_val > 5:
            raise ValueError("Level must be beetween 0 and 5")

        def update_action_values(actions, new_val, old_vals=None):
            diff = bool(old_vals)
            if not old_vals:
                old_vals = defaultdict(dict)
            for k, v in actions['ClearSkies'].items():
                if not diff:
                    old_vals['ClearSkies'][k] = int(v[0]['@Value'].value)
                mod = new_val - old_vals['ClearSkies'][k]
                for i in v:
                    i['@Value'].value = str((int(i['@Value'].value) + mod) if diff else new_val)
            for action in ('Complete', 'ItemSelected', 'Started'):
                for k, v in actions[action].items():
                    if not diff:
                        old_vals[action][k] = int(v[0]['@Value'].value)
                    mod = (new_val * 2) - old_vals[action][k]
                    for i in v:
                        i['@Value'].value = str((int(i['@Value'].value) + mod)
                                                if diff else (new_val * 2))
            if not diff:
                old_vals['Complete']['All'] = int(actions['Complete']['All'][0]['@Value'].value)
            mod = (new_val * 2 * (len(Pony.GameTypes.map) + 1)) - old_vals['Complete']['All']
            for i in actions['Complete']['All']:
                i['@Value'].value = str((int(i['@Value'].value) + mod)
                                        if diff else (new_val * 2 *
                                                      (len(Pony.GameTypes.map) + 1)))
            return old_vals

        old_vals = update_action_values(self._pony_actions['Pony'], new_val)
        update_action_values(self._pony_actions['Global'], new_val, old_vals)
        self.shards = 0
        self._leveltag["@Level"].value = str(new_val)

    @StoredItem.shards.setter
    def shards(self, new_val):
        new_val = int(new_val)
        if new_val < 0 or new_val > 10:
            raise ValueError("Shards must be beetween 0 and 10")
        if self.level == 5 and new_val > 0:
            raise ValueError("Can't have shards at maximum level")
        self._leveltag["@Shards"].value = str(new_val)

    @StoredItem.next_game.getter
    def next_game(self):
        try:
            gametype = int(self._minigametag["@NextPlayAction"].value)
            name = Pony.GameTypes.rmap[gametype]
        except:
            name = 'Unknown'
        return name

    @next_game.setter
    def next_game(self, gametype):
        if gametype not in Pony.GameTypes.rmap:
            raise ValueError("Invalid game type")
        self._minigametag["@NextPlayAction"].value = str(gametype)

    def level_up(self, nb=1):
        self.level = min(max(self.level + nb, 0), 5)
        self.shards = 0

    def level_down(self, nb=1):
        self.level_up(-nb)

    def shard_up(self):
        if self.level != 5:
            self.shards = 10

    def shard_down(self):
        if self.level != 5:
            self.shards = 0

    def reset_game_timer(self):
        self._minigametag["@NextPlayTime"].value = ZERO def retrieve_gluid(db_path):
    try:
        db = sqlite3.connect(db_path)
        cur = db.execute('SELECT value FROM glshare WHERE key = ?',
                         ('ANMP.GloftPOHM_GAIA_ENC_KEY_GLUID',))
        value, = cur.fetchone()
        res = json.loads(value)
        gluid = res['data']
    except:
        raise GluidError('Unable to retrieve GLUID from file')
    return gluid def check_version(force=False):
    res = {
        "up_to_date": True,
        "error": None,
        "download_url": None,
    }
    if not check_network():
        res["error"] = "It seems you don't have access to the internet"
    elif not check_frozen() and not force:
        res["error"] = "You're not running the frozen version"
    else:
        r = requests.get('https://api.github.com/repos/Arzaroth/CelestiaSunrise/releases/latest')
        res["up_to_date"] = parse_version(r.json()['tag_name']) <= parse_version('.'.join(VERSION))
        try:
            res["download_url"] = r.json()['assets'][0]['browser_download_url']
        except:
            res["error"] = "Unable to find download url"
    return res def docopt_cmd(func):
    @wraps(func)
    def wrapper(self, args):
        from docopt import docopt, DocoptExit
        try:
            args = docopt(func.__doc__, args)
        except DocoptExit as e:
            print(e)
            return
        except SystemExit:
            return
        return func(self, args)
    return wrapper def docopt_cmd_completion(func, **kwargs):
    options = parse_defaults(func.__doc__)
    pattern = parse_pattern(formal_usage(printable_usage(func.__doc__)),
                            options).children[0]

    def get_state(it, pattern):
        try:
            value = next(it)
        except StopIteration:
            return pattern
        res = []
        for x in pattern:
            if ((type(x[0]) == list and value in flatten(x[0])) or
                value == x[0]):
                res.append(x[1:])
        if res:
            return get_state(it, res)
        return []

    def wrapper(self, text, line, begidx, endidx):
        argv = shlex.split(line[:endidx])[1:]
        if not line[endidx - 1].isspace():
            target = argv[-1]
            argv = argv[:-1]
        else:
            target = ''
        state = get_state(iter(argv), pattern.noflat())
        res = []
        for x in state:
            if type(x[0]) == list:
                res.extend(flatten(x[0]))
            else:
                res.append(x[0])
        return list(set(x for x in res if x.startswith(target)))

    wrapper.__name__ = str('complete_' + func.__name__[3:])
    wrapper.__module__ = func.__module__
    wrapper.__doc__ = func.__doc__
    return wrapper def show_player(xml_handle, args):
    for typ in xml_handle.player_infos.values():
        if not args['<player_data>'] or args['<player_data>'] == typ.name:
            print(typ) def show_currencies(xml_handle, args):
    for name, typ in xml_handle.currencies.items():
        print('{}:'.format(name))
        for cur, val in typ.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) def show_ponies(xml_handle, args):
    if args['-o']:
        print('Not owned ponies:')
        for ID, name in xml_handle.missing_ponies.items():
            print('Pony(ID: {}, Name: {})'.format(ID, name))
    else:
        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[0])) 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][0]) 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 show_quests(xml_handle, args):
    print('Quests:')
    for quest in xml_handle.quests:
        print('  {}'.format(quest)) def set_player(xml_handle, args):
    for typ in xml_handle.player_infos.values():
        if args['<player_data>'] == typ.name:
            try:
                typ.value = args['<value>']
            except ValueError as e:
                print(str(e)) 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 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 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) def set_inventory(xml_handle, args):
    for pony_id in args['<not_owned_pony_id>']:
        if pony_id in xml_handle.missing_ponies:
            xml_handle.inventory.append(pony_id)
            xml_handle.missing_ponies.remove(pony_id) def set_quests(xml_handle, args):
    if args['complete']:
        for quest in xml_handle.quests:
            quest.complete()
import binascii
import sys
import six
from cmd import Cmd
from celestia.save import SaveManager, SaveError
from celestia.save import decompress_data, compress_data
from celestia.xml.xmlhandler import XmlHandler
from celestia.utility.utility import Pony
from celestia.utility.gluid import retrieve_gluid
from celestia.utility.version import check_version
from .docopt_utils import docopt_cmd, docopt_cmd_completion
from .show import (show_player,
                   show_currencies, show_currency,
                   show_ponies, show_pony,
                   show_zones, show_zone,
                   show_quests)
from .set import (set_player, set_currency,
                  set_ponies, set_pony,
                  set_zones, set_zone,
                  set_inventory, set_quests)
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# File: ponyshell.py
# by Arzaroth Lekva
# lekva@arzaroth.com
#

from __future__ import print_function, absolute_import, unicode_literals



class PonyMeta(type):
    def __new__(cls, name, bases, attrs):
        new_attrs = {}
        for key, value in attrs.items():
            if key.startswith('do_') and value.__doc__:
                new_attrs['complete_' + key[3:]] = docopt_cmd_completion(value)
            new_attrs[key] = value
        return super(PonyMeta, cls).__new__(cls, name, bases, new_attrs)


@six.add_metaclass(PonyMeta)
class PonyShell(Cmd, object):
    prompt = 'ponyshell> '

    def __init__(self, savedata):
        Cmd.__init__(self)
        self._save_manager = SaveManager(savedata.savefile, savedata.gluid)
        data, self.save_number = self._save_manager.load(savedata.legacy)
        if not savedata.legacy:
            data = decompress_data(data)
        self._xml_handle = XmlHandler(data)
        self._xml_handle.pre_load()
        self.savedata = savedata
        self._show_functions = {
            'player': show_player,
            'currencies': show_currencies,
            'currency': show_currency,
            'ponies': show_ponies,
            'pony': show_pony,
            'zones': show_zones,
            'zone': show_zone,
            'quests': show_quests,
        }
        self._set_functions = {

            'player': set_player,