From ad765659dd6e50fbd78ad921af05dc66ceadb0e3 Mon Sep 17 00:00:00 2001
From: CaitSith2 <d_good@caitsith2.com>
Date: Tue, 20 Oct 2020 01:16:20 -0700
Subject: [PATCH] Add Race rom encryption.

---
 Main.py          |  2 +-
 Rom.py           | 36 ++++++++++++++++++++++++++++--------
 requirements.txt |  3 ++-
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/Main.py b/Main.py
index 75d2bf19..d58e0c00 100644
--- a/Main.py
+++ b/Main.py
@@ -188,7 +188,7 @@ def main(args, seed=None):
             patch_enemizer(world, player, rom, args.enemizercli)
 
         if args.race:
-            patch_race_rom(rom)
+            patch_race_rom(rom, world, player)
 
         world.spoiler.hashes[(player, team)] = get_hash_string(rom.hash)
 
diff --git a/Rom.py b/Rom.py
index 0fe66a5d..f3acb99d 100644
--- a/Rom.py
+++ b/Rom.py
@@ -11,6 +11,7 @@ import struct
 import sys
 import subprocess
 import threading
+import xxtea
 import concurrent.futures
 from typing import Optional
 
@@ -58,6 +59,31 @@ class LocalRom(object):
     def write_bytes(self, startaddress: int, values):
         self.buffer[startaddress:startaddress + len(values)] = values
 
+    def encrypt_range(self, startaddress: int, length: int, key: bytes):
+        for i in range(0, length, 8):
+            data = bytes(self.read_bytes(startaddress + i, 8))
+            data = xxtea.encrypt(data, key, padding=False)
+            self.write_bytes(startaddress + i, bytearray(data))
+
+    def encrypt(self, world, player):
+        local_random = world.rom_seeds[player]
+        key = bytes(local_random.getrandbits(8 * 16).to_bytes(16, 'big'))
+        self.write_bytes(0x1800B0, bytearray(key))
+        self.write_int16(0x180087, 1)
+
+        itemtable = []
+        locationtable = []
+        for i in range(168):
+            itemtable.append(self.read_byte(0xE96E + (i * 3)))
+            locationtable.append(self.read_byte(0xe96C + (i * 3)))
+            locationtable.append(self.read_byte(0xe96D + (i * 3)))
+        self.write_bytes(0xE96C, locationtable)
+        self.write_bytes(0xE96C + 0x150, itemtable)
+        self.encrypt_range(0xE96C + 0x150, 168, key)
+        self.encrypt_range(0x180000, 32, key)
+        self.encrypt_range(0x180140, 32, key)
+
+
     def write_to_file(self, file, hide_enemizer=False):
         with open(file, 'wb') as outfile:
             outfile.write(self.buffer)
@@ -1422,16 +1448,10 @@ def patch_rom(world, rom, player, team, enemized):
 
     return rom
 
-try:
-    import RaceRom
-except ImportError:
-    RaceRom = None
 
-def patch_race_rom(rom):
+def patch_race_rom(rom, world, player):
     rom.write_bytes(0x180213, [0x01, 0x00]) # Tournament Seed
-
-    if 'RaceRom' in sys.modules:
-        RaceRom.encrypt(rom)
+    rom.encrypt(world, player)
 
 def write_custom_shops(rom, world, player):
     shops = [shop for shop in world.shops if shop.custom and shop.region.player == player]
diff --git a/requirements.txt b/requirements.txt
index 24e7fc59..7efc2049 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,4 +5,5 @@ fuzzywuzzy>=0.18.0
 bsdiff4>=1.2.0
 prompt_toolkit>=3.0.6
 appdirs>=1.4.4
-maseya-z3pr>=1.0.0rc1
\ No newline at end of file
+maseya-z3pr>=1.0.0rc1
+xxtea>=2.0.0
\ No newline at end of file