fix temp file deletion when scan is already running

This commit is contained in:
2024-06-16 00:26:33 +02:00
parent d507b84266
commit a5d5448001

View File

@@ -25,8 +25,6 @@ import struct
import sys import sys
import time import time
from subprocess import Popen, PIPE
uvscan_regex = re.compile(r"Found:?(?: the| potentially unwanted program| (?:virus|trojan) or variant)? (.+?)(?:\.| (?:virus |trojan )?)", re.MULTILINE) uvscan_regex = re.compile(r"Found:?(?: the| potentially unwanted program| (?:virus|trojan) or variant)? (.+?)(?:\.| (?:virus |trojan )?)", re.MULTILINE)
@@ -61,6 +59,7 @@ class AIO(asyncio.Protocol):
raise RuntimeError("queue not set") raise RuntimeError("queue not set")
self.logger = logging.getLogger(__name__) self.logger = logging.getLogger(__name__)
self.tmpfile = None self.tmpfile = None
self.cancelled = False
def _send_response(self, response): def _send_response(self, response):
response = response.encode() + AIO.separator response = response.encode() + AIO.separator
@@ -68,7 +67,6 @@ class AIO(asyncio.Protocol):
self.transport.write(response) self.transport.write(response)
self.transport.close() self.transport.close()
def connection_made(self, transport): def connection_made(self, transport):
self.peer = transport.get_extra_info("peername") self.peer = transport.get_extra_info("peername")
self.logger.info("new connection from {}".format(self.peer)) self.logger.info("new connection from {}".format(self.peer))
@@ -79,7 +77,6 @@ class AIO(asyncio.Protocol):
self.command = None self.command = None
self.length = None self.length = None
self.all_chunks = False self.all_chunks = False
self.completed = False
def data_received(self, data): def data_received(self, data):
try: try:
@@ -135,39 +132,48 @@ class AIO(asyncio.Protocol):
self._send_response(str(e)) self._send_response(str(e))
def process_uvscan_result(self, result): def process_uvscan_result(self, result):
self.logger.info("{} received uvscan result of {}: {}".format(self.peer, self.tmpfile, result)) self.logger.debug("{} removing temporary file {}".format(self.peer, self.tmpfile))
self.completed = True os.remove(self.tmpfile)
self._send_response(result) self.tmpfile = None
if not self.cancelled:
self.logger.info("{} received uvscan result of {}: {}".format(self.peer, self.tmpfile, result))
self._send_response(result)
def connection_lost(self, exc): def connection_lost(self, exc):
if self.tmpfile: if self.tmpfile:
if not self.completed: entries = []
self.logger.warning("{} client prematurely closed connection, removing {} from scan queue".format(self.peer, self.tmpfile)) try:
entries = [] for entry in iter(AIO.queue.get_nowait, None):
try: if not entry:
for entry in iter(AIO.queue.get_nowait, None): continue
if not entry: if entry[1] != self.tmpfile:
continue entries.append(entry)
if entry[1] != self.tmpfile: else:
entries.append(entry) self.cancelled = True
except asyncio.QueueEmpty: except asyncio.QueueEmpty:
pass pass
for entry in entries: for entry in entries:
AIO.queue.put_nowait(entry) AIO.queue.put_nowait(entry)
self.logger.debug("{} removing temporary file {}".format(self.peer, self.tmpfile)) if self.cancelled:
os.remove(self.tmpfile) self.logger.warning("{} client prematurely closed connection, skipped scan of {}".format(self.peer, self.tmpfile))
self.logger.info("closed connection to {}".format(self.peer)) self.logger.debug("{} removing temporary file {}".format(self.peer, self.tmpfile))
os.remove(self.tmpfile)
else:
self.logger.warning("{} client prematurely closed connection".format(self.peer))
self.cancelled = True
else:
self.logger.info("closed connection to {}".format(self.peer))
def main(): def main():
"Run uvscand." "Run uvscand."
# parse command line # parse command line
parser = argparse.ArgumentParser(description="uvscand daemon", parser = argparse.ArgumentParser(
description="uvscand daemon",
formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=45, width=140)) formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=45, width=140))
parser.add_argument("-c", "--config", help="List of config files to read.", nargs="+", parser.add_argument("-c", "--config", help="List of config files to read.", nargs="+", default=["/etc/uvscand.conf"])
default=["/etc/uvscand.conf"]) parser.add_argument("-m", "--maxprocs", help="Maximum number of parallel scan processes.", type=int, default=8)
parser.add_argument("-m", "--maxprocs", help="Maximum number of parallel scan processes.",
type=int, default=8)
parser.add_argument("-d", "--debug", help="Log debugging messages.", action="store_true") parser.add_argument("-d", "--debug", help="Log debugging messages.", action="store_true")
args = parser.parse_args() args = parser.parse_args()