962377f6ebcbe96bcfee795379f440f72e0d374d
[zzz-pokedex.git] / 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.ItemInternalID) \
132 .filter_by(generation_id = 4, internal_id = 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.ItemInternalID) \
177 .filter_by(generation_id = 4, internal_id = 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.LocationInternalID) \
185 .filter_by(generation_id = 4, internal_id = egg_loc_id).one().location
186
187 self._met_location = session.query(tables.LocationInternalID) \
188 .filter_by(generation_id = 4, internal_id = 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