pokedex/struct/__init__.py

   1 # encoding: utf8
   2 u"""
   3 Handles reading and encryption/decryption of Pokémon save file data.
   4
   5 See: http://projectpokemon.org/wiki/Pokemon_NDS_Structure
   6
   7 Kudos to LordLandon for his pkmlib.py, from which this module was originally
   8 derived.
   9 """
  10
  11 import struct
  12
  13 from pokedex.db import tables
  14 from pokedex.formulae import calculated_hp, calculated_stat
  15 from pokedex.util import namedtuple, permutations
  16 from pokedex.struct._pokemon_struct import pokemon_struct
  17
  18 def pokemon_prng(seed):
  19     u"""Creates a generator that simulates the main Pokémon PRNG."""
  20     while True:
  21         seed = 0x41C64E6D * seed + 0x6073
  22         seed &= 0xFFFFFFFF
  23         yield seed >> 16
  24
  25
  26 class SaveFilePokemon(object):
  27     u"""Represents an individual Pokémon, from the game's point of view.
  28
  29     Handles translating between the on-disk encrypted form, the in-RAM blob
  30     (also used by pokesav), and something vaguely intelligible.
  31     """
  32
  33     Stat = namedtuple('Stat', ['stat', 'base', 'gene', 'exp', 'calc'])
  34
  35     def __init__(self, blob, encrypted=False):
  36         u"""Wraps a Pokémon save struct in a friendly object.
  37
  38         If `encrypted` is True, the blob will be decrypted as though it were an
  39         on-disk save.  Otherwise, the blob is taken to be already decrypted and
  40         is left alone.
  41
  42         `session` is an optional database session.
  43         """
  44
  45         if encrypted:
  46             # Decrypt it.
  47             # Interpret as one word (pid), followed by a bunch of shorts
  48             struct_def = "I" + "H" * ((len(blob) - 4) / 2)
  49             shuffled = list( struct.unpack(struct_def, blob) )
  50
  51             # Apply standard Pokémon decryption, undo the block shuffling, and
  52             # done
  53             self.reciprocal_crypt(shuffled)
  54             words = self.shuffle_chunks(shuffled, reverse=True)
  55             self.blob = struct.pack(struct_def, *words)
  56
  57         else:
  58             # Already decrypted
  59             self.blob = blob
  60
  61         self.structure = pokemon_struct.parse(self.blob)
  62
  63     @property
  64     def as_struct(self):
  65         u"""Returns a decrypted struct, aka .pkm file."""
  66         return self.blob
  67
  68     @property
  69     def as_encrypted(self):
  70         u"""Returns an encrypted struct the game expects in a save file."""
  71
  72         # Interpret as one word (pid), followed by a bunch of shorts
  73         struct_def = "I" + "H" * ((len(self.blob) - 4) / 2)
  74         words = list( struct.unpack(struct_def, self.blob) )
  75
  76         # Apply the block shuffle and standard Pokémon encryption
  77         shuffled = self.shuffle_chunks(words)
  78         self.reciprocal_crypt(shuffled)
  79
  80         # Stuff back into a string, and done
  81         return struct.pack(struct_def, *shuffled)
  82
  83     ### Delicious data
  84     @property
  85     def is_shiny(self):
  86         u"""Returns true iff this Pokémon is shiny."""
  87         # See http://bulbapedia.bulbagarden.net/wiki/Personality#Shininess
  88         # But don't see it too much, because the above is super over
  89         # complicated.  Do this instead!
  90         personality_msdw = self.structure.personality >> 16
  91         personality_lsdw = self.structure.personality & 0xffff
  92         return (
  93             self.structure.original_trainer_id
  94             ^ self.structure.original_trainer_secret_id
  95             ^ personality_msdw
  96             ^ personality_lsdw
  97         ) < 8
  98
  99     def use_database_session(self, session):
 100         """Remembers the given database session, and prefetches a bunch of
 101         database stuff.  Gotta call this before you use the database properties
 102         like `species`, etc.
 103         """
 104         self._session = session
 105
 106         st = self.structure
 107         self._pokemon = session.query(tables.Pokemon).get(st.national_id)
 108         self._pokemon_form = session.query(tables.PokemonForm) \
 109             .with_parent(self._pokemon) \
 110             .filter_by(name=st.alternate_form) \
 111             .one()
 112         self._ability = self._session.query(tables.Ability).get(st.ability_id)
 113
 114         growth_rate = self._pokemon.evolution_chain.growth_rate
 115         self._experience_rung = session.query(tables.Experience) \
 116             .filter(tables.Experience.growth_rate == growth_rate) \
 117             .filter(tables.Experience.experience <= st.exp) \
 118             .order_by(tables.Experience.level.desc()) \
 119             [0]
 120         level = self._experience_rung.level
 121
 122         self._next_experience_rung = None
 123         if level < 100:
 124             self._next_experience_rung = session.query(tables.Experience) \
 125                 .filter(tables.Experience.growth_rate == growth_rate) \
 126                 .filter(tables.Experience.level == level + 1) \
 127                 .one()
 128
 129         self._held_item = None
 130         if st.held_item_id:
 131             self._held_item = session.query(tables.ItemGameIndex) \
 132                 .filter_by(generation_id = 4, game_index = st.held_item_id).one().item
 133
 134         self._stats = []
 135         for pokemon_stat in self._pokemon.stats:
 136             structure_name = pokemon_stat.stat.name.lower().replace(' ', '_')
 137             gene = st.ivs['iv_' + structure_name]
 138             exp  = st['effort_' + structure_name]
 139
 140             if pokemon_stat.stat.name == u'HP':
 141                 calc = calculated_hp
 142             else:
 143                 calc = calculated_stat
 144
 145             stat_tup = self.Stat(
 146                 stat = pokemon_stat.stat,
 147                 base = pokemon_stat.base_stat,
 148                 gene = gene,
 149                 exp  = exp,
 150                 calc = calc(
 151                     pokemon_stat.base_stat,
 152                     level = level,
 153                     iv = gene,
 154                     effort = exp,
 155                 ),
 156             )
 157
 158             self._stats.append(stat_tup)
 159
 160
 161         move_ids = (
 162             self.structure.move1_id,
 163             self.structure.move2_id,
 164             self.structure.move3_id,
 165             self.structure.move4_id,
 166         )
 167         move_rows = self._session.query(tables.Move).filter(tables.Move.id.in_(move_ids))
 168         moves_dict = dict((move.id, move) for move in move_rows)
 169
 170         self._moves = [moves_dict.get(move_id, None) for move_id in move_ids]
 171
 172         if st.hgss_pokeball >= 17:
 173             pokeball_id = st.hgss_pokeball - 17 + 492
 174         else:
 175             pokeball_id = st.dppt_pokeball
 176         self._pokeball = session.query(tables.ItemGameIndex) \
 177             .filter_by(generation_id = 4, game_index = pokeball_id).one().item
 178
 179         egg_loc_id = st.pt_egg_location_id or st.dp_egg_location_id
 180         met_loc_id = st.pt_met_location_id or st.dp_met_location_id
 181
 182         self._egg_location = None
 183         if egg_loc_id:
 184             self._egg_location = session.query(tables.LocationGameIndex) \
 185                 .filter_by(generation_id = 4, game_index = egg_loc_id).one().location
 186
 187         self._met_location = session.query(tables.LocationGameIndex) \
 188             .filter_by(generation_id = 4, game_index = met_loc_id).one().location
 189
 190     @property
 191     def species(self):
 192         # XXX forme!
 193         return self._pokemon
 194
 195     @property
 196     def species_form(self):
 197         return self._pokemon_form
 198
 199     @property
 200     def pokeball(self):
 201         return self._pokeball
 202
 203     @property
 204     def egg_location(self):
 205         return self._egg_location
 206
 207     @property
 208     def met_location(self):
 209         return self._met_location
 210
 211     @property
 212     def shiny_leaves(self):
 213         return (
 214             self.structure.shining_leaves.leaf1,
 215             self.structure.shining_leaves.leaf2,
 216             self.structure.shining_leaves.leaf3,
 217             self.structure.shining_leaves.leaf4,
 218             self.structure.shining_leaves.leaf5,
 219         )
 220
 221     @property
 222     def level(self):
 223         return self._experience_rung.level
 224
 225     @property
 226     def exp_to_next(self):
 227         if self._next_experience_rung:
 228             return self._next_experience_rung.experience - self.structure.exp
 229         else:
 230             return 0
 231
 232     @property
 233     def progress_to_next(self):
 234         if self._next_experience_rung:
 235             return 1.0 \
 236                 * (self.structure.exp - self._experience_rung.experience) \
 237                 / (self._next_experience_rung.experience - self._experience_rung.experience)
 238         else:
 239             return 0.0
 240
 241     @property
 242     def ability(self):
 243         return self._ability
 244
 245     @property
 246     def held_item(self):
 247         return self._held_item
 248
 249     @property
 250     def stats(self):
 251         return self._stats
 252
 253     @property
 254     def moves(self):
 255         return self._moves
 256
 257     @property
 258     def move_pp(self):
 259         return (
 260             self.structure.move1_pp,
 261             self.structure.move2_pp,
 262             self.structure.move3_pp,
 263             self.structure.move4_pp,
 264         )
 265
 266
 267     ### Utility methods
 268
 269     shuffle_orders = list( permutations(range(4)) )
 270
 271     @classmethod
 272     def shuffle_chunks(cls, words, reverse=False):
 273         """The main 128 encrypted bytes (or 64 words) in a save block are split
 274         into four chunks and shuffled around in some order, based on
 275         personality.  The actual order of shuffling is a permutation of four
 276         items in order, indexed by the shuffle index.  That is, 0 yields 0123,
 277         1 yields 0132, 2 yields 0213, etc.
 278
 279         Given a list of words (the first of which should be the pid), this
 280         function returns the words in shuffled order.  Pass reverse=True to
 281         unshuffle instead.
 282         """
 283
 284         pid = words[0]
 285         shuffle_index = (pid >> 0xD & 0x1F) % 24
 286
 287         shuffle_order = cls.shuffle_orders[shuffle_index]
 288         if reverse:
 289             # Decoding requires going the other way; invert the order
 290             shuffle_order = [shuffle_order.index(i) for i in range(4)]
 291
 292         shuffled = words[:3]  # skip the unencrypted stuff
 293         for chunk in shuffle_order:
 294             shuffled += words[ chunk * 16 + 3 : chunk * 16 + 19 ]
 295         shuffled += words[67:]  # extra bytes are also left alone
 296
 297         return shuffled
 298
 299     @classmethod
 300     def reciprocal_crypt(cls, words):
 301         u"""Applies the reciprocal Pokémon save file cipher to the provided
 302         list of words.
 303
 304         Returns nothing; the list is changed in-place.
 305         """
 306         # Apply regular Pokémon "encryption": xor everything with the output of
 307         # the PRNG.  First three items are pid/unused/checksum and are not
 308         # encrypted.
 309
 310         # Main data is encrypted using the checksum as a seed
 311         prng = pokemon_prng(words[2])
 312         for i in range(3, 67):
 313             words[i] ^= next(prng)
 314
 315         if len(words) > 67:
 316             # Extra bytes are encrypted using the pid as a seed
 317             prng = pokemon_prng(words[0])
 318             for i in range(67, len(words)):
 319                 words[i] ^= next(prng)
 320
 321         return