# HG changeset patch
# User Thibaut Girka <thib@sitedethib.com>
# Date 1325268649 -3600
# Node ID 1c24a6d93c1b5a5b733b2af60a289eb244d7993c
# Parent  5afc75f71fed186d228ceee1b9439e81e10a458a
Improve find_character_defs, clean up a bit, and disable attack types 2 and 3 for now

diff --git a/pytouhou/formats/exe.py b/pytouhou/formats/exe.py
--- a/pytouhou/formats/exe.py
+++ b/pytouhou/formats/exe.py
@@ -34,13 +34,10 @@ class Shot(object):
         self.speed = 0.
         self.damage = 0
         self.orb = 0
-        self.shot_type = 0
+        self.type = 0
         self.sprite = 0
         self.unknown1 = None
 
-    def __repr__(self):
-        return '(%d, %d, %f, %f, %d, %d, %d, %d, %d)' % (self.interval, self.delay, self.angle, self.speed, self.damage, self.orb, self.shot_type, self.sprite, self.unknown1)
-
 
 class SHT(object):
     def __init__(self):
@@ -61,33 +58,63 @@ class SHT(object):
 
 
     @classmethod
-    def find_character_records(self, file, pe_file):
+    def find_character_defs(cls, pe_file):
+        """Generator returning the possible VA of character definition blocks.
+
+        Based on knowledge of the structure, it tries to find valid definition blocks
+        without embedding any copyrighted material or hard-coded offsets that would
+        only be useful for a specific build of the game.
+        """
+
         format = Struct('<4f2I')
-        data_section = [section for section in pe_file.sections if section.Name.startswith('.data')][0]
-        text_section = [section for section in pe_file.sections if section.Name.startswith('.text')][0]
+        data_section = [section for section in pe_file.sections
+                            if section.Name.startswith('.data')][0]
+        text_section = [section for section in pe_file.sections
+                            if section.Name.startswith('.text')][0]
         data_va = pe_file.image_base + data_section.VirtualAddress
+        data_size = data_section.SizeOfRawData
         text_va = pe_file.image_base + text_section.VirtualAddress
+        text_size = text_section.SizeOfRawData
 
-        for addr in xrange(data_va, data_va + data_section.SizeOfRawData, 4):
+        # Search the whole data segment for 4 successive character definitions
+        for addr in xrange(data_va, data_va + data_size, 4):
             for character_id in xrange(4):
                 pe_file.seek_to_va(addr + character_id * 24)
-                speed1, speed2, speed3, speed4, ptr1, ptr2 = format.unpack(file.read(format.size))
+                (speed1, speed2, speed3, speed4,
+                 ptr1, ptr2) = format.unpack(pe_file.file.read(format.size))
 
-                if not (all(0. < x < 8. for x in (speed1, speed2, speed3, speed4))
+                # Check whether the character's speed make sense,
+                # and whether the function pointers point to valid addresses
+                if not (all(0. < x < 10. for x in (speed1, speed2, speed3, speed4))
                         and speed2 <= speed1
-                        and 0 <= ptr1 - text_va < text_section.SizeOfRawData - 8
-                        and 0 <= ptr2 - text_va < text_section.SizeOfRawData - 8):
+                        and 0 <= ptr1 - text_va < text_size - 8
+                        and 0 <= ptr2 - text_va < text_size - 8):
                     break
 
+                # So far, this character definition seems to be valid.
+                # Now, make sure the shoot function wrappers pass valid addresses
                 pe_file.seek_to_va(ptr1 + 4)
-                shtptr1, = unpack('<I', file.read(4))
+                shtptr1, = unpack('<I', pe_file.file.read(4))
                 pe_file.seek_to_va(ptr2 + 4)
-                shtptr2, = unpack('<I', file.read(4))
+                shtptr2, = unpack('<I', pe_file.file.read(4))
+                if not (0 <= shtptr1 - data_va < data_size - 12
+                        and 0 <= shtptr2 - data_va < data_size - 12):
+                    break
 
-                if not (0 <= shtptr1 - data_va < data_section.SizeOfRawData
-                        and 0 <= shtptr2 - data_va < data_section.SizeOfRawData):
+                # It is unlikely this character record is *not* valid, but
+                # just to be sure, let's check the first SHT definition.
+                pe_file.seek_to_va(shtptr1)
+                nb_shots, power, shotsptr = unpack('<III', pe_file.file.read(12))
+                if not (0 < nb_shots <= 1000
+                        and 0 <= power < 1000
+                        and 0 <= shotsptr - data_va < data_size - 36*nb_shots):
                     break
-            else: # XXX: Obscure python feature! This gets executed if there were no break!
+
+            else:
+                # XXX: Obscure python feature! This only gets executed if the
+                # XXX: loop ended without a break statement.
+                # In our case, it's only executed if all the 4 character
+                # definitions are considered valid.
                 yield addr
 
 
@@ -95,7 +122,7 @@ class SHT(object):
     def read(cls, file):
         pe_file = PEFile(file)
 
-        character_records_va = list(cls.find_character_records(file, pe_file))[0]
+        character_records_va = list(cls.find_character_defs(pe_file))[0]
 
         characters = []
         shots_offsets = []
@@ -142,7 +169,7 @@ class SHT(object):
 
                     data = unpack('<HH6fHBBhh', file.read(36))
                     (shot.interval, shot.delay, x, y, hitbox_x, hitbox_y,
-                     shot.angle, shot.speed, shot.damage, shot.orb, shot.shot_type,
+                     shot.angle, shot.speed, shot.damage, shot.orb, shot.type,
                      shot.sprite, shot.unknown1) = data
 
                     shot.pos = (x, y)
diff --git a/pytouhou/games/eosd.py b/pytouhou/games/eosd.py
--- a/pytouhou/games/eosd.py
+++ b/pytouhou/games/eosd.py
@@ -51,7 +51,6 @@ class EoSDGame(Game):
                       ItemType(etama3, 6, 13)] #Star
 
         characters = resource_loader.get_eosd_characters('102h.exe')
-        print characters[0]
 
         #eosd_characters = [ReimuA, ReimuB, MarisaA, MarisaB]
         players = []
@@ -135,6 +134,9 @@ class EoSDPlayer(Player):
         nb_bullets_max = self._game.nb_bullets_max
 
         for shot in sht.shots[power]:
+            if shot.type in (2, 3): # TODO: Those shot types aren't handled yet
+                continue
+
             if self.fire_time % shot.interval == shot.delay:
                 if nb_bullets_max is not None and len(bullets) == nb_bullets_max:
                     break
@@ -145,9 +147,11 @@ class EoSDPlayer(Player):
 
                 #TODO: find a better way to do that.
                 bullet_type = BulletType(self.anm_wrapper, shot.sprite % 256,
-                                                         shot.sprite % 256 + 32, #TODO: find the real cancel anim
-                                                         0, 0, 0, 0.)
+                                         shot.sprite % 256 + 32, #TODO: find the real cancel anim
+                                         0, 0, 0, 0.)
                 bullets.append(Bullet((x, y), bullet_type, 0,
                                       shot.angle, shot.speed,
                                       (0, 0, 0, 0, 0., 0., 0., 0.),
-                                      0, self, self._game, player_bullet=True, damage=shot.damage, hitbox=shot.hitbox))
+                                      0, self, self._game, player_bullet=True,
+                                      damage=shot.damage, hitbox=shot.hitbox))
+