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