pokedex/savefile.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 from pokedex.util import permutations
  13
  14 def pokemon_prng(seed):
  15     u"""Creates a generator that simulates the main Pokémon PRNG."""
  16     while True:
  17         seed = 0x41C64E6D * seed + 0x6073
  18         seed &= 0xFFFFFFFF
  19         yield seed >> 16
  20
  21
  22 class PokemonSave(object):
  23     u"""Represents an individual Pokémon, from the game's point of view.
  24
  25     Handles translating between the on-disk encrypted form, the in-RAM blob
  26     (also used by pokesav), and something vaguely intelligible.
  27
  28     XXX: Okay, well, right now it's just encryption and decryption.  But, you
  29     know.
  30     """
  31
  32     def __init__(self, blob, encrypted=False):
  33         u"""Wraps a Pokémon save struct in a friendly object.
  34
  35         If `encrypted` is True, the blob will be decrypted as though it were an
  36         on-disk save.  Otherwise, the blob is taken to be already decrypted and
  37         is left alone.
  38         """
  39
  40         # XXX Sometime this should have an abstract internal representation.
  41         # For now, just store the decrypted version
  42
  43         if encrypted:
  44             # Decrypt it.
  45             # Interpret as one word (pid), followed by a bunch of shorts
  46             struct_def = "I" + "H" * ((len(blob) - 4) / 2)
  47             shuffled = list( struct.unpack(struct_def, blob) )
  48
  49             # Apply standard Pokémon decryption, undo the block shuffling, and
  50             # done
  51             self.reciprocal_crypt(shuffled)
  52             words = self.shuffle_chunks(shuffled, reverse=True)
  53             self.blob = struct.pack(struct_def, *words)
  54
  55         else:
  56             # Already decrypted
  57             self.blob = blob
  58
  59
  60     @property
  61     def as_struct(self):
  62         u"""Returns a decrypted struct, aka .pkm file."""
  63         return self.blob
  64
  65     @property
  66     def as_encrypted(self):
  67         u"""Returns an encrypted struct the game expects in a save file."""
  68
  69         # Interpret as one word (pid), followed by a bunch of shorts
  70         struct_def = "I" + "H" * ((len(self.blob) - 4) / 2)
  71         words = list( struct.unpack(struct_def, self.blob) )
  72
  73         # Apply the block shuffle and standard Pokémon encryption
  74         shuffled = self.shuffle_chunks(words)
  75         self.reciprocal_crypt(shuffled)
  76
  77         # Stuff back into a string, and done
  78         return struct.pack(struct_def, *shuffled)
  79
  80
  81     ### Utility methods
  82
  83     shuffle_orders = list( permutations(range(4)) )
  84
  85     @classmethod
  86     def shuffle_chunks(cls, words, reverse=False):
  87         """The main 128 encrypted bytes (or 64 words) in a save block are split
  88         into four chunks and shuffled around in some order, based on
  89         personality.  The actual order of shuffling is a permutation of four
  90         items in order, indexed by the shuffle index.  That is, 0 yields 0123,
  91         1 yields 0132, 2 yields 0213, etc.
  92
  93         Given a list of words (the first of which should be the pid), this
  94         function returns the words in shuffled order.  Pass reverse=True to
  95         unshuffle instead.
  96         """
  97
  98         pid = words[0]
  99         shuffle_index = (pid >> 0xD & 0x1F) % 24
 100
 101         shuffle_order = cls.shuffle_orders[shuffle_index]
 102         if reverse:
 103             # Decoding requires going the other way; invert the order
 104             shuffle_order = [shuffle_order.index(i) for i in range(4)]
 105
 106         shuffled = words[:3]  # skip the unencrypted stuff
 107         for chunk in shuffle_order:
 108             shuffled += words[ chunk * 16 + 3 : chunk * 16 + 19 ]
 109         shuffled += words[67:]  # extra bytes are also left alone
 110
 111         return shuffled
 112
 113     @classmethod
 114     def reciprocal_crypt(cls, words):
 115         u"""Applies the reciprocal Pokémon save file cipher to the provided
 116         list of words.
 117
 118         Returns nothing; the list is changed in-place.
 119         """
 120         # Apply regular Pokémon "encryption": xor everything with the output of
 121         # the PRNG.  First three items are pid/unused/checksum and are not
 122         # encrypted.
 123
 124         # Main data is encrypted using the checksum as a seed
 125         prng = pokemon_prng(words[2])
 126         for i in range(3, 67):
 127             words[i] ^= next(prng)
 128
 129         if len(words) > 67:
 130             # Extra bytes are encrypted using the pid as a seed
 131             prng = pokemon_prng(words[0])
 132             for i in range(67, len(words)):
 133                 words[i] ^= next(prng)
 134
 135         return