Fixes dockerfile issues + better temp file handling

Author: pufi656

docker/Dockerfile.zeratool_lib

@@ -1,25 +1,28 @@
-FROM python:3.10
+FROM python:3.11
 
 # Install and configure Poetry
 RUN pip3 install poetry
 
 # Setup the convertor from pyproject.toml to requirements.txt
-RUN pip3 install click toml
+RUN pip3 install click toml tomli
 RUN wget https://raw.githubusercontent.com/jla524/requirements/main/requirements/convert.py -O /convert.py
 
 # Convert Zeratool's dependencies
 COPY /docker/zeratool_lib /zeratool_lib
-RUN python3 /convert.py --noversion /zeratool_lib
+RUN python3 /convert.py --no-version /zeratool_lib
 RUN mv /zeratool_lib/requirements.txt /requirements.zeratool_lib.txt
 
 # Copies service's dependencies
 COPY /docker/requirements.txt /requirements.service.txt
 
 # Convert module's dependencies
 COPY /pyproject.toml /pyproject.toml
-RUN python3 /convert.py --noversion /
+RUN python3 /convert.py --no-version /
 RUN mv /requirements.txt /requirements.aeg.txt
 
+# Download the commons library
+RUN git clone https://github.com/CyberReasoningSystem/commons /commons
+
 # Install all dependencies
 RUN pip3 install -r /requirements.service.txt
 RUN pip3 install -r /requirements.aeg.txt
@@ -30,9 +33,6 @@ COPY /automatic_exploit_generation /automatic_exploit_generation
 COPY /docker/protobuf /protobuf
 COPY /docker/zeratool_lib_service.py /zeratool_lib_service.py
 
-# Download the commons library
-RUN git clone https://github.com/CyberReasoningSystem/commons /commons
-
 # Install Radare2
 RUN set -xe; \
     wget https://github.com/radareorg/radare2/releases/download/5.9.0/radare2_5.9.0_amd64.deb; \

docker/zeratool_lib_service.py

@@ -15,42 +15,49 @@
 YEAR_IN_SECONDS = 365 * 60 * 60 * 24
 
 
-def _create_temp_binary(content: bytes) -> tempfile.NamedTemporaryFile:
-    binary = tempfile.NamedTemporaryFile()
-
-    binary.write(content)
-    binary.flush()
-
-    os.chmod(binary.name, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)
-
-    return binary
+def _create_temp_binary(content: bytes) -> str:
+    binary = tempfile.NamedTemporaryFile(delete=False)
+    try:
+        binary.write(content)
+        binary.flush()
+        path = binary.name
+    finally:
+        binary.close()
+    os.chmod(path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)
+    return path
 
 
 class EService(exploit_pb2_grpc.ExploitServiceServicer):
     def Exploit(self, request, _):
-        temp_file = _create_temp_binary(request.binary)
-
-        input_stream = ZeratoolInputStreams(request.input_stream)
-        overflow_only = request.overflow_only
-        format_only = request.format_only
-        win_funcs_used = request.serialized_win_funcs.split(",")
-        if len(win_funcs_used) == 1 and win_funcs_used[0] == '':
-            win_funcs_used = None
-
-        print("temp_file.name: ", temp_file.name)
-        print("overflow_only: ", overflow_only)
-        print("format_only: ", format_only)
-        print("win_funcs_used: ", win_funcs_used)
-
-        result = exploit(
-            temp_file.name,
-            input_stream,
-            overflow_only=overflow_only,
-            format_only=format_only,
-            win_funcs=win_funcs_used,
-        )
-
-        return exploit_pb2.exploit(pickledExploit=pickle.dumps(result))
+        temp_path = _create_temp_binary(request.binary)
+        try:
+            input_stream = ZeratoolInputStreams(request.input_stream)
+            overflow_only = request.overflow_only
+            format_only = request.format_only
+            win_funcs_used = request.serialized_win_funcs.split(",")
+            if len(win_funcs_used) == 1 and win_funcs_used[0] == '':
+                win_funcs_used = None
+
+            print("temp_path: ", temp_path)
+            print("overflow_only: ", overflow_only)
+            print("format_only: ", format_only)
+            print("win_funcs_used: ", win_funcs_used)
+
+            result = exploit(
+                temp_path,
+                input_stream,
+                overflow_only=overflow_only,
+                format_only=format_only,
+                win_funcs=win_funcs_used,
+            )
+
+            return exploit_pb2.exploit(pickledExploit=pickle.dumps(result))
+        finally:
+            if os.path.exists(temp_path):
+                try:
+                    os.unlink(temp_path)
+                except OSError:
+                    pass
 
 
 def serve() -> None:

pyproject.toml

@@ -23,7 +23,7 @@ r2pipe = "^1.8.0"
 timeout-decorator = "^0.5.0"
 tqdm = "^4.65.0"
 commons = {path = "../commons"}
-zeratool_lib = {path = "../zeratool_lib"}
+zeratool_lib = {path = "../zeratool_lib", develop = true}
 
 [tool.poetry.dev-dependencies]
 black = "^22.6.0"