diff --git a/d2warehouse/fileformat.py b/d2warehouse/fileformat.py new file mode 100644 index 0000000..23369bd --- /dev/null +++ b/d2warehouse/fileformat.py @@ -0,0 +1,19 @@ +# Copyright 2023 +# Copyright 2023 +# +# This file is part of d2warehouse. +# +# d2warehouse is free software: you can redistribute it and/or modify it under the +# terms of the GNU General Public License as published by the Free Software +# Foundation, either version 3 of the License, or (at your option) any later +# version. +# +# d2warehouse is distributed in the hope that it will be useful, but WITHOUT ANY +# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR +# A PARTICULAR PURPOSE. See the GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License along with +# Mercator. If not, see . + +STASH_TAB_MAGIC = b"\x55\xAA\x55\xAA" +ITEM_DATA_MAGIC = b"JM" diff --git a/d2warehouse/item.py b/d2warehouse/item.py index ffb241e..e1908f3 100644 --- a/d2warehouse/item.py +++ b/d2warehouse/item.py @@ -19,6 +19,7 @@ import os import re from typing import Optional from bitarray import bitarray +from bitarray.util import int2ba from dataclasses import dataclass from enum import Enum @@ -181,9 +182,25 @@ class Item: bits = bitarray(endian="little") bits.frombytes(self.raw_data) print(" " * indent, txtbits(bits)) + print(" " * indent, self.raw_data.hex()) print("") print("") + def set_position(self, x: int, y: int) -> None: + if x < 0 or x > 15 or y < 0 or y > 15: + raise ValueError("position needs to be in range [0, 16)") + self.pos_x = x + self.pos_y = y + bits = bitarray(endian="little") + bits.frombytes(self.raw_data) + bits[42:46] = int2ba(x, 4, endian="little") + bits[46:50] = int2ba(y, 4, endian="little") + self.raw_data = bits.tobytes() + + def size(self) -> tuple[int, int]: + base = lookup_basetype(self.code) + return base["width"], base["height"] + def lookup_basetype(code: str) -> dict: global _basetype_map diff --git a/d2warehouse/parser.py b/d2warehouse/parser.py index 3059359..67f10d2 100644 --- a/d2warehouse/parser.py +++ b/d2warehouse/parser.py @@ -27,9 +27,7 @@ from d2warehouse.item import ( lookup_stat, ) import d2warehouse.huffman as huffman - -STASH_TAB_MAGIC = b"\x55\xAA\x55\xAA" -ITEM_DATA_MAGIC = b"JM" +from d2warehouse.fileformat import STASH_TAB_MAGIC, ITEM_DATA_MAGIC class ParseError(RuntimeError): diff --git a/d2warehouse/stash.py b/d2warehouse/stash.py index ff04afd..fc24e35 100644 --- a/d2warehouse/stash.py +++ b/d2warehouse/stash.py @@ -15,15 +15,102 @@ # You should have received a copy of the GNU General Public License along with # Mercator. If not, see . from d2warehouse.item import Item +from d2warehouse.fileformat import STASH_TAB_MAGIC, ITEM_DATA_MAGIC +from struct import pack + + +class StashError(RuntimeError): + pass + + +class StashFullError(StashError): + pass class Stash: def __init__(self) -> None: self.tabs: list[StashTab] = [] + def raw(self) -> bytes: + return b"".join(tab.raw() for tab in self.tabs) + class StashTab: def __init__(self) -> None: self.gold: int = 0 + self.version: int = 99 self.item_data: bytes = b"" self.items: list[Item] = [] + + def raw(self) -> bytes: + """Get the computed raw representation of the stash""" + item_raw = b"".join(item.raw_data for item in self.items) + raw_length = len(item_raw) + 0x44 + return ( + STASH_TAB_MAGIC + + pack("I", 1) + + pack("I", self.version) + + pack("I", self.gold) + + pack("I", raw_length) + + b"\x00" * 44 + + ITEM_DATA_MAGIC + + pack("H", len(self.items)) + + item_raw + ) + + def remove(self, item: Item) -> None: + for idx, current_item in enumerate(self.items): + if item is current_item: + self.items.pop(idx) + return + raise StashError("Item not found in stash") + + def add(self, item: Item) -> None: + width, height = item.size() + dest = self._find_destination(width, height) + if not dest: + raise StashFullError( + "Could not locate an open spot in the stash to add the item" + ) + x, y = dest + item.set_position(x, y) + self.items.append(item) + + @staticmethod + def _is_free_location( + slots: list[list[bool]], x: int, y: int, width: int, height: int + ) -> bool: + """Check if a given location can fit an item of a given width and height""" + if x + width > len(slots): + return False + if y + height > len(slots[0]): + return False + for i in range(width): + for j in range(height): + if slots[x + i][y + j]: + return False + return True + + def _find_destination(self, width: int, height: int) -> tuple[int, int] | None: + """Find a free destination to store an item of given width and height""" + slots = self._build_slot_matrix() + stash_width = len(slots) + stash_height = len(slots[0]) + for y in range(stash_width): + for x in range(stash_height): + if self._is_free_location(slots, x, y, width, height): + return x, y + return None + + def _build_slot_matrix(self, rows: int = 10, cols: int = 10) -> list[list[bool]]: + """Build a matrix of filled and empty slots for the stash tab.""" + slots = [] + for i in range(cols): + slots.append([False] * rows) + for item in self.items: + x, y = item.pos_x, item.pos_y + width, height = item.size() + for n in range(width): + for m in range(height): + slots[x + n][y + m] = True + return slots diff --git a/d2warehouse/tests/test_stash.py b/d2warehouse/tests/test_stash.py new file mode 100644 index 0000000..6fffb1c --- /dev/null +++ b/d2warehouse/tests/test_stash.py @@ -0,0 +1,86 @@ +import unittest +from d2warehouse.parser import parse_stash, parse_item +from d2warehouse.stash import StashFullError + + +class StashTest(unittest.TestCase): + @staticmethod + def simple_stash(): + # tab 0 contains a rune, a buckler and arrows + # tab 1 contains a gem + # tab 3 contains nothing + # All tabs contain some gold + data = bytes.fromhex( + "55aa55aa0100000063000000a40100007400000000000000000000000000" + "000000000000000000000000000000000000000000000000000000000000" + "000000004a4d03001000a0000500f47cfb0010008000050454a1083f57a9" + "190c01041818fc0710008000050cf46deec23225b20c4340ff3f55aa55aa" + "0100000063000000391b00004d0000000000000000000000000000000000" + "000000000000000000000000000000000000000000000000000000000000" + "4a4d01001000a000050474891555aa55aa01000000630000003905000044" + "000000000000000000000000000000000000000000000000000000000000" + "00000000000000000000000000000000004a4d0000" + ) + stash = parse_stash(data) + return data, stash + + def dupe(self, item): + data = item.raw_data + data, duped = parse_item(data) + self.assertEqual(data, b"") + return duped + + def test_raw_equals(self): + data, stash = self.simple_stash() + self.assertEqual(len(stash.tabs), 3) + self.assertEqual([len(tab.items) for tab in stash.tabs], [3, 1, 0]) + gold = tuple(tab.gold for tab in stash.tabs) + self.assertEqual(gold, (420, 6969, 1337)) + self.assertEqual(data, stash.raw()) + + def test_remove(self): + data, stash = self.simple_stash() + rune, buckler, arrows = stash.tabs[0].items[:3] + gem = stash.tabs[1].items[0] + + stash.tabs[0].remove(buckler) + self.assertEqual(stash.tabs[0].items, [rune, arrows]) + stash.tabs[0].remove(rune) + self.assertEqual(stash.tabs[0].items, [arrows]) + stash.tabs[0].remove(arrows) + self.assertEqual(stash.tabs[0].items, []) + stash.tabs[1].remove(gem) + self.assertEqual(stash.tabs[1].items, []) + + new_stash = parse_stash(stash.raw()) + self.assertEqual(len(new_stash.tabs[0].items), 0) + self.assertEqual(len(new_stash.tabs[1].items), 0) + self.assertEqual(len(new_stash.tabs[2].items), 0) + + def test_add(self): + data, stash = self.simple_stash() + rune, buckler, arrows = stash.tabs[0].items[:3] + stash.tabs[0].remove(buckler) + stash.tabs[1].add(buckler) + self.assertEqual(buckler.pos_x, 2) + self.assertEqual(buckler.pos_y, 0) + + new_stash = parse_stash(stash.raw()) + self.assertEqual(len(new_stash.tabs[0].items), 2) + self.assertEqual(len(new_stash.tabs[1].items), 2) + + def test_overflow(self): + data, stash = self.simple_stash() + rune, buckler, arrows = stash.tabs[0].items[:3] + bucklers = [self.dupe(buckler) for _ in range(25)] + + for buckler in bucklers: + stash.tabs[2].add(buckler) + + with self.assertRaises(StashFullError): + stash.tabs[2].add(self.dupe(buckler)) + + new_stash = parse_stash(stash.raw()) + self.assertEqual(len(new_stash.tabs[0].items), 3) + self.assertEqual(len(new_stash.tabs[1].items), 1) + self.assertEqual(len(new_stash.tabs[2].items), 25)