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     ### Delicious data
  80
  81     @property
  82     def is_shiny(self):
  83         u"""Returns true iff this Pokémon is shiny."""
  84         # See http://bulbapedia.bulbagarden.net/wiki/Personality#Shininess
  85         # But don't see it too much, because the above is super over
  86         # complicated.  Do this instead!
  87         personality_msdw = self.structure.personality >> 16
  88         personality_lsdw = self.structure.personality & 0xffff
  89         return (
  90             self.structure.original_trainer_id
  91             ^ self.structure.original_trainer_secret_id
  92             ^ personality_msdw
  93             ^ personality_lsdw
  94         ) < 8
  95
  96
  97     ### Utility methods
  98
  99     shuffle_orders = list( permutations(range(4)) )
 100
 101     @classmethod
 102     def shuffle_chunks(cls, words, reverse=False):
 103         """The main 128 encrypted bytes (or 64 words) in a save block are split
 104         into four chunks and shuffled around in some order, based on
 105         personality.  The actual order of shuffling is a permutation of four
 106         items in order, indexed by the shuffle index.  That is, 0 yields 0123,
 107         1 yields 0132, 2 yields 0213, etc.
 108
 109         Given a list of words (the first of which should be the pid), this
 110         function returns the words in shuffled order.  Pass reverse=True to
 111         unshuffle instead.
 112         """
 113
 114         pid = words[0]
 115         shuffle_index = (pid >> 0xD & 0x1F) % 24
 116
 117         shuffle_order = cls.shuffle_orders[shuffle_index]
 118         if reverse:
 119             # Decoding requires going the other way; invert the order
 120             shuffle_order = [shuffle_order.index(i) for i in range(4)]
 121
 122         shuffled = words[:3]  # skip the unencrypted stuff
 123         for chunk in shuffle_order:
 124             shuffled += words[ chunk * 16 + 3 : chunk * 16 + 19 ]
 125         shuffled += words[67:]  # extra bytes are also left alone
 126
 127         return shuffled
 128
 129     @classmethod
 130     def reciprocal_crypt(cls, words):
 131         u"""Applies the reciprocal Pokémon save file cipher to the provided
 132         list of words.
 133
 134         Returns nothing; the list is changed in-place.
 135         """
 136         # Apply regular Pokémon "encryption": xor everything with the output of
 137         # the PRNG.  First three items are pid/unused/checksum and are not
 138         # encrypted.
 139
 140         # Main data is encrypted using the checksum as a seed
 141         prng = pokemon_prng(words[2])
 142         for i in range(3, 67):
 143             words[i] ^= next(prng)
 144
 145         if len(words) > 67:
 146             # Extra bytes are encrypted using the pid as a seed
 147             prng = pokemon_prng(words[0])
 148             for i in range(67, len(words)):
 149                 words[i] ^= next(prng)
 150
 151         return