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