fix temp file deletion when scan is already running
This commit is contained in:
@@ -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()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user