diff --git a/src/clamd/__init__.py b/src/clamd/__init__.py index 92ff640..91c01ce 100644 --- a/src/clamd/__init__.py +++ b/src/clamd/__init__.py @@ -63,8 +63,11 @@ def _init_socket(self): """ try: self.clamd_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.clamd_socket.connect((self.host, self.port)) + # 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) + self.clamd_socket.connect((self.host, self.port)) except socket.error: e = sys.exc_info()[1] @@ -166,11 +169,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")} @@ -184,15 +189,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() @@ -231,7 +234,7 @@ def _send_command(self, cmd, *args): concat_args = ' ' + ' '.join(args) cmd = 'n{cmd}{args}\n'.format(cmd=cmd, args=concat_args).encode('utf-8') - self.clamd_socket.send(cmd) + self.clamd_socket.sendall(cmd) def _recv_response(self): """