From 182faaf3c68d0c229316cc742661ddaeb0df31de Mon Sep 17 00:00:00 2001
From: Thomas Oettli <spacefreak@noop.ch>
Date: Wed, 29 Jan 2020 23:39:46 +0100
Subject: [PATCH] Try to fix left over tmp files

Signed-off-by: Thomas Oettli <spacefreak@noop.ch>
---
 setup.py            |  2 +-
 uvscand/__init__.py | 16 ++++++++++------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/setup.py b/setup.py
index 7eaea4b..482de14 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ def read_file(fname):
         return f.read()
 
 setup(name = "uvscand",
-    version = "0.0.1",
+    version = "0.0.2",
     author = "Thomas Oettli",
     author_email = "spacefreak@noop.ch",
     description = "A python daemon to perform virus scans with uvscan (McAfee) over TCP socket.",
diff --git a/uvscand/__init__.py b/uvscand/__init__.py
index 0b26b5a..e867a09 100755
--- a/uvscand/__init__.py
+++ b/uvscand/__init__.py
@@ -34,7 +34,6 @@ uvscan_regex = re.compile(r"Found:?(?: the| potentially unwanted program| (?:vir
 async def run(uvscan, filename):
     proc = await asyncio.create_subprocess_exec(uvscan, "--secure", "--mime", "--noboot", "--panalyse", "--manalyse", filename, stdout=asyncio.subprocess.PIPE)
     stdout, _ = await proc.communicate()
-    os.remove(filename)
     if proc.returncode == 13:
         match = uvscan_regex.search(stdout.decode())
         name = match.group(1) if match else "UNKNOWN"
@@ -53,6 +52,7 @@ class AIO(asyncio.Protocol):
             raise RuntimeError("configuration not set")
         self.logger = logging.getLogger(__name__)
         self.data = bytearray()
+        self.tmpfile = None
 
     def connection_made(self, transport):
         self.peer = transport.get_extra_info("peername")
@@ -79,17 +79,17 @@ class AIO(asyncio.Protocol):
             pos += 1
             if command == "zINSTREAM":
                 # save data chunks to temporary file
-                tmpfile = os.path.join(AIO.config["tmpdir"], "uvscan_{}_{}".format(self.request_time, str(self.peer[1])))
-                self.logger.debug("save data from {} in temporary file {}".format(self.peer, tmpfile))
-                with open(tmpfile, "wb") as f:
+                self.tmpfile = os.path.join(AIO.config["tmpdir"], "uvscan_{}_{}".format(self.request_time, str(self.peer[1])))
+                self.logger.debug("save data from {} in temporary file {}".format(self.peer, self.tmpfile))
+                with open(self.tmpfile, "wb") as f:
                     while True:
                         length = struct.unpack(">I", self.data[pos:pos + 4])[0]
                         if length == 0: break
                         pos += 4
                         f.write(self.data[pos:pos + length])
                         pos += length
-                self.logger.debug("starting uvscan for file {}".format(tmpfile))
-                task = asyncio.async(run(AIO.config["uvscan_path"], tmpfile))
+                self.logger.debug("starting uvscan for file {}".format(self.tmpfile))
+                task = asyncio.async(run(AIO.config["uvscan_path"], self.tmpfile))
                 task.add_done_callback(self.handle_uvscan_result)
             else:
                 raise RuntimeError("unknown command")
@@ -106,6 +106,10 @@ class AIO(asyncio.Protocol):
         self.transport.write(response)
         self.transport.close()
 
+    def connection_list(self, exc):
+        if self.tmpfile:
+            os.remove(self.tmpfile)
+
 
 def main():
     "Run uvscand."