Skip to content

Commit 86b58b4

Browse files
committed
apply changes from graingert#23
- use sendall() instead of send() - allow custom max_buffer_size - set socket timeout before connect Thanks to @mrname
1 parent b0b5225 commit 86b58b4

File tree

1 file changed

+11
-8
lines changed

1 file changed

+11
-8
lines changed

clammy/__init__.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,13 +36,16 @@ def _init_socket(self):
3636
try:
3737
self.clamd_socket = socket.socket(self.socket_type, socket.SOCK_STREAM)
3838

39+
# Set timeout prior to connecting to ensure that an initial
40+
# connection timeout will respect the setting regardless of OS.
41+
# https://docs.python.org/3/library/socket.html#timeouts-and-the-connect-method
42+
self.clamd_socket.settimeout(self.timeout)
43+
3944
if self.socket_type == socket.AF_INET:
4045
self.clamd_socket.connect((self.host, self.port))
4146
elif self.socket_type == socket.AF_UNIX:
4247
self.clamd_socket.connect(self.unix_socket)
4348

44-
self.clamd_socket.settimeout(self.timeout)
45-
4649
except socket.error:
4750
if self.socket_type == socket.AF_UNIX:
4851
error_message = f'Error connecting to Unix socket "{self.unix_socket}"'
@@ -129,11 +132,13 @@ def _file_system_scan(self, command, file):
129132
finally:
130133
self._close_socket()
131134

132-
def instream(self, buff):
135+
def instream(self, buff, max_chunk_size=1024):
133136
"""
134137
Scan a buffer
135138
136139
buff filelikeobj: buffer to scan
140+
max_chunk_size int: Maximum size of chunk to send to clamd in bytes
141+
MUST be < StreamMaxLength in /etc/clamav/clamd.conf
137142
138143
return:
139144
- (dict): {filename1: ("virusname", "status")}
@@ -147,15 +152,13 @@ def instream(self, buff):
147152
self._init_socket()
148153
self._send_command("INSTREAM")
149154

150-
max_chunk_size = 1024 # MUST be < StreamMaxLength in /etc/clamav/clamd.conf
151-
152155
chunk = buff.read(max_chunk_size)
153156
while chunk:
154157
size = struct.pack(b"!L", len(chunk))
155-
self.clamd_socket.send(size + chunk)
158+
self.clamd_socket.sendall(size + chunk)
156159
chunk = buff.read(max_chunk_size)
157160

158-
self.clamd_socket.send(struct.pack(b"!L", 0))
161+
self.clamd_socket.sendall(struct.pack(b"!L", 0))
159162

160163
result = self._recv_response()
161164

@@ -195,7 +198,7 @@ def _send_command(self, cmd, *args):
195198

196199
# cmd = 'n{cmd}{args}\n'.format(cmd=cmd, args=concat_args).encode('utf-8')
197200
cmd = f"n{cmd}{concat_args}\n".encode("utf-8")
198-
self.clamd_socket.send(cmd)
201+
self.clamd_socket.sendall(cmd)
199202

200203
def _recv_response(self):
201204
"""

0 commit comments

Comments
 (0)