diff --git a/clammy/__init__.py b/clammy/__init__.py index cb5bdd2..fec19f1 100644 --- a/clammy/__init__.py +++ b/clammy/__init__.py @@ -36,13 +36,16 @@ def _init_socket(self): try: self.clamd_socket = socket.socket(self.socket_type, socket.SOCK_STREAM) + # Set timeout prior to connecting to ensure that an initial + # connection timeout will respect the setting regardless of OS. + # https://docs.python.org/3/library/socket.html#timeouts-and-the-connect-method + self.clamd_socket.settimeout(self.timeout) + if self.socket_type == socket.AF_INET: self.clamd_socket.connect((self.host, self.port)) elif self.socket_type == socket.AF_UNIX: self.clamd_socket.connect(self.unix_socket) - self.clamd_socket.settimeout(self.timeout) - except socket.error: if self.socket_type == socket.AF_UNIX: error_message = f'Error connecting to Unix socket "{self.unix_socket}"' @@ -129,11 +132,13 @@ def _file_system_scan(self, command, file): finally: self._close_socket() - def instream(self, buff): + def instream(self, buff, max_chunk_size=1024): """ Scan a buffer buff filelikeobj: buffer to scan + max_chunk_size int: Maximum size of chunk to send to clamd in bytes + MUST be < StreamMaxLength in /etc/clamav/clamd.conf return: - (dict): {filename1: ("virusname", "status")} @@ -147,15 +152,13 @@ def instream(self, buff): self._init_socket() self._send_command("INSTREAM") - max_chunk_size = 1024 # MUST be < StreamMaxLength in /etc/clamav/clamd.conf - chunk = buff.read(max_chunk_size) while chunk: size = struct.pack(b"!L", len(chunk)) - self.clamd_socket.send(size + chunk) + self.clamd_socket.sendall(size + chunk) chunk = buff.read(max_chunk_size) - self.clamd_socket.send(struct.pack(b"!L", 0)) + self.clamd_socket.sendall(struct.pack(b"!L", 0)) result = self._recv_response() @@ -195,7 +198,7 @@ def _send_command(self, cmd, *args): # cmd = 'n{cmd}{args}\n'.format(cmd=cmd, args=concat_args).encode('utf-8') cmd = f"n{cmd}{concat_args}\n".encode("utf-8") - self.clamd_socket.send(cmd) + self.clamd_socket.sendall(cmd) def _recv_response(self): """