Merge pull request #14 from Phrohdoh/swgb

Parse SWGB DRS archives
fredreichbier · Sep 24, 2015 · 069b884 · 069b884
2 parents a5edcbc + b11cdea
commit 069b884
Showing 1 changed file with 19 additions and 2 deletions.
diff --git a/genie/drs.py b/genie/drs.py
@@ -42,7 +42,7 @@ def _decode(self, obj, context):
     )
 )
 
-HEADER = cons.Struct('header',
+AOE_HEADER = cons.Struct('header',
     cons.String('copyright', 40, padchar='\0'),
     cons.String('version', 4),
     cons.String('file_type', 12, padchar='\0'),
@@ -51,6 +51,15 @@ def _decode(self, obj, context):
     cons.Array(lambda ctx: ctx['number_of_tables'], TableAdapter(TABLE)),
 )
 
+SWGB_HEADER = cons.Struct(
+    cons.String('copyright', 60, padchar='\0'),
+    cons.String('version', 4),
+    cons.String('file_type', 12, padchar='\0'),
+    cons.ULInt32('number_of_tables'),
+    cons.ULInt32('offset'),
+    cons.Array(lambda ctx: ctx['number_of_tables'], TableAdapter(TABLE)),
+)
+
 def get_file_extension(resource_type):
     """
         get the embedded file extension from the genie resource type (a 4-byte number).
@@ -106,7 +115,15 @@ class DRSFile(object):
     """
     def __init__(self, stream):
         self.stream = stream
-        self.header = HEADER._parse(stream, cons.Container(drs_file=self))
+        pos = stream.tell()
+        stream.seek(64)
+        maybe_swgb_header = stream.read(4)
+        stream.seek(pos)
+
+        if maybe_swgb_header == "swbg":
+            self.header = SWGB_HEADER._parse(stream, cons.Container(drs_file=self))
+        else:
+            self.header = AOE_HEADER._parse(stream, cons.Container(drs_file=self))
 
     @property
     def tables(self):