From 0ed859b95e4bdc25c3247d0f949b7c8861249fa1 Mon Sep 17 00:00:00 2001
From: hakril <project@hakril.net>
Date: Thu, 31 Oct 2019 16:27:01 +0100
Subject: [PATCH] Added support for simple expression in array (like [MAX_PATH
 + 1])

---
 ctypes_generation/generate.py      |  5 +++--
 ctypes_generation/simpleparser.py  |  6 ++++-
 ctypes_generation/struct_parser.py | 12 +++++++---
 ctypes_generation/winstruct.py     | 36 +++++++++++++++++++++++++-----
 4 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/ctypes_generation/generate.py b/ctypes_generation/generate.py
index 050f1fe..f92f074 100644
--- a/ctypes_generation/generate.py
+++ b/ctypes_generation/generate.py
@@ -87,8 +87,9 @@ class StructureParsedFile(ParsedFile):
                 self.imports_by_struct[field_type.name] = struct.name
             try:
                 int(nb_rep)
-            except:
-                self.add_imports(nb_rep)
+            except TypeError:
+                # nb_rep is a ComplexArrayExpression
+                self.add_imports(*nb_rep.get_names())
 
     def asser_struct_not_already_in_import(self, struct):
         for sname in [struct.name] + struct.typedef.keys():
diff --git a/ctypes_generation/simpleparser.py b/ctypes_generation/simpleparser.py
index da4cf35..b842942 100644
--- a/ctypes_generation/simpleparser.py
+++ b/ctypes_generation/simpleparser.py
@@ -39,6 +39,10 @@ class StarToken(NoValueToken):
     value = "*"
     pass
 
+class PlusToken(NoValueToken):
+    value = "+"
+    pass
+
 class OpenBracketToken(NoValueToken):
     value = "{"
     pass
@@ -79,7 +83,7 @@ class Lexer(object):
     token_chr = {"*" : StarToken, "[" : OpenSquareBracketToken, "]" : CloseSquareBracketToken,
                     "{" : OpenBracketToken, "}" : CloseBracketToken, ";" : SemiColonToken,
                     "," : CommaToken, "(" : OpenParenthesisToken, ")" :  CloseParenthesisToken, "#" : SharpToken, "=" : EqualToken,
-                    ":": ColonToken}
+                    ":": ColonToken, "+": PlusToken}
 
     def __init__(self, code, newlinetoken=False):
         self.code = code
diff --git a/ctypes_generation/struct_parser.py b/ctypes_generation/struct_parser.py
index f40606d..1a4be4b 100644
--- a/ctypes_generation/struct_parser.py
+++ b/ctypes_generation/struct_parser.py
@@ -1,6 +1,6 @@
 import dummy_wintypes
 import itertools
-from winstruct import WinStruct, WinUnion, WinStructType, Ptr, WinEnum, BitFieldValue
+from winstruct import WinStruct, WinUnion, WinStructType, Ptr, WinEnum, BitFieldValue, ComplexArrayExpression
 from simpleparser import *
 
 
@@ -14,9 +14,15 @@ class WinStructParser(Parser):
         if type(self.peek()) == OpenSquareBracketToken:
             # Array
             self.assert_token_type(OpenSquareBracketToken)
-            number = self.assert_token_type(NameToken).value
+            nb_rep = ComplexArrayExpression()
+            while type(self.peek()) != CloseSquareBracketToken:
+                tok = self.next_token()
+                if type(tok) not in (NameToken, PlusToken):
+                    raise ValueError("Array expression only accept names/+")
+                is_name = (type(tok) == NameToken)
+                nb_rep.add_token_to_expression(tok.value, is_name)
             self.assert_token_type(CloseSquareBracketToken)
-            return number
+            return nb_rep # YOLO, import
         # Bitfield
         self.assert_token_type(ColonToken)
         nb_bits = self.promote_to_int(self.next_token())
diff --git a/ctypes_generation/winstruct.py b/ctypes_generation/winstruct.py
index 5ff815f..adbba72 100644
--- a/ctypes_generation/winstruct.py
+++ b/ctypes_generation/winstruct.py
@@ -14,6 +14,30 @@ class BitFieldValue(object):
     def __int__(self):
         return self.nb_bits
 
+    def generate_ctypes(self):
+        return self.nb_bits
+
+class ComplexArrayExpression(object):
+    def __init__(self):
+        self.values = []
+        self.names = []
+
+    def add_token_to_expression(self, value, isname=False):
+        self.values.append(value)
+        if isname:
+            try:
+                int(value) # An INT ? ignore
+            except ValueError:
+                self.names.append(value)
+
+    def get_names(self):
+        return self.names
+
+    def generate_ctypes(self):
+        # if len(self.values) > 1:
+            # import pdb;pdb.set_trace()
+        return "({0})".format(" ".join(self.values))
+
 class WinStructType(object):
     def __init__(self, name):
         self.name = name
@@ -107,9 +131,9 @@ class WinStruct(object):
             res += ["{0}._pack_ = ".format(self.pack)]
         res += ["{0}._fields_ = [".format(self.name)]
         for (ftype, name, nb_rep) in self.fields:
-            if isinstance(nb_rep, BitFieldValue):
-                res += ['    ("{0}", {1}, {2}),'.format(name, ftype.generate_ctypes(), nb_rep.nb_bits)]
-            elif  nb_rep == 1:
+            if isinstance(nb_rep, (BitFieldValue, ComplexArrayExpression)):
+                nb_rep = nb_rep.generate_ctypes()
+            if nb_rep == 1:
                 res+= ['    ("{0}", {1}),'.format(name, ftype.generate_ctypes())]
             else:
                 res+= ['    ("{0}", {1} * {2}),'.format(name, ftype.generate_ctypes(), nb_rep)]
@@ -135,10 +159,10 @@ class WinStruct(object):
 
 
         for (ftype, name, nb_rep) in self.fields:
-            if isinstance(nb_rep, BitFieldValue):
+            if isinstance(nb_rep, (BitFieldValue, ComplexArrayExpression)):
                 # Should I check 'ftype' somewhere when we have a bitfield ?
-                res+= '        ("{0}", {1}, {2}),\n'.format(name, ftype.generate_ctypes(), nb_rep.nb_bits)
-            elif  nb_rep == 1:
+                nb_rep = nb_rep.generate_ctypes()
+            if nb_rep == 1:
                 res+= '        ("{0}", {1}),\n'.format(name, ftype.generate_ctypes())
             else:
                 res+= '        ("{0}", {1} * {2}),\n'.format(name, ftype.generate_ctypes(), nb_rep)