5
5
import os
6
6
import re
7
7
import ssl
8
- import sys
9
8
import time
10
- import urllib .parse
11
9
import uuid
12
10
13
11
# Configuration
34
32
num_submitted = 0
35
33
num_processed = 0
36
34
35
+ # URL to use for submission
36
+ api_endpoint = ""
37
+
38
+ # Original args
39
+ args = {}
37
40
38
41
# Functions
39
42
def process_dir (workdir ):
@@ -106,53 +109,70 @@ def skip_file(filepath):
106
109
def submit_sample (filepath ):
107
110
print ("[SUBMIT] Submitting {} ..." .format (filepath ))
108
111
112
+ headers = {
113
+ "Content-Type" : "application/octet-stream" ,
114
+ "Content-Disposition" : f"attachment; filename={ filepath } " ,
115
+ }
116
+
109
117
try :
110
- headers = {
111
- "Content-Type" : "application/octet-stream" ,
112
- "Content-Disposition" : f"attachment; filename={ filepath } " ,
113
- }
114
118
115
119
with open (filepath , "rb" ) as f :
116
120
data = f .read ()
117
121
118
- boundary = str (uuid .uuid4 ())
119
- headers = {
120
- "Content-Type" : f"multipart/form-data; boundary={ boundary } " ,
121
- }
122
-
123
- # Create multipart/form-data payload
124
- payload = (
125
- f"--{ boundary } \r \n "
126
- f'Content-Disposition: form-data; name="file"; filename="{ filepath } "\r \n '
127
- f"Content-Type: application/octet-stream\r \n \r \n "
128
- ).encode ("utf-8" )
129
- payload += data
130
- payload += f"\r \n --{ boundary } --\r \n " .encode ("utf-8" )
131
-
132
- if args .tls :
133
- if args .insecure :
134
- context = ssl ._create_unverified_context ()
122
+ except Exception as e :
123
+ print ("[ERROR] Could not read '{}' - {}" .format (filepath , e ))
124
+ return
125
+
126
+ boundary = str (uuid .uuid4 ())
127
+ headers = {
128
+ "Content-Type" : f"multipart/form-data; boundary={ boundary } " ,
129
+ }
130
+
131
+ # Create multipart/form-data payload
132
+ payload = (
133
+ f"--{ boundary } \r \n "
134
+ f'Content-Disposition: form-data; name="file"; filename="{ filepath } "\r \n '
135
+ f"Content-Type: application/octet-stream\r \n \r \n "
136
+ ).encode ("utf-8" )
137
+ payload += data
138
+ payload += f"\r \n --{ boundary } --\r \n " .encode ("utf-8" )
139
+
140
+ retries = 0
141
+ while retries < 3 :
142
+ try :
143
+ if args .tls :
144
+ if args .insecure :
145
+ context = ssl ._create_unverified_context ()
146
+ else :
147
+ context = ssl ._create_default_https_context ()
148
+ conn = http .client .HTTPSConnection (args .server , args .port , context = context )
135
149
else :
136
- context = ssl ._create_default_https_context ()
137
- conn = http .client .HTTPSConnection (args .server , args .port , context = context )
138
- else :
139
- conn = http .client .HTTPConnection (args .server , args .port )
140
- conn .request ("POST" , api_endpoint , body = payload , headers = headers )
150
+ conn = http .client .HTTPConnection (args .server , args .port )
151
+ conn .request ("POST" , api_endpoint , body = payload , headers = headers )
141
152
142
- resp = conn .getresponse ()
153
+ resp = conn .getresponse ()
143
154
144
- global num_submitted
145
- num_submitted += 1
155
+ except Exception as e :
156
+ print ("[ERROR] Could not submit '{}' - {}" .format (filepath , e ))
157
+ retries += 1
158
+ time .sleep (2 << retries )
159
+ continue
146
160
147
- if resp .status != 200 :
148
- print (
149
- "[ERROR] HTTP return status: {}, reason: {}" .format (
150
- resp .status , resp .reason
151
- )
161
+ # pylint: disable=no-else-continue
162
+ if resp .status == 503 : # Service unavailable
163
+ retry_time = resp .headers .get ("Retry-After" , 30 )
164
+ time .sleep (retry_time )
165
+ continue
166
+ elif resp .status == 200 :
167
+ break
168
+ print (
169
+ "[ERROR] HTTP return status: {}, reason: {}" .format (
170
+ resp .status , resp .reason
152
171
)
172
+ )
153
173
154
- except Exception as e :
155
- print ( "Could not submit '{}' - {}" . format ( filepath , e ))
174
+ global num_submitted
175
+ num_submitted += 1
156
176
157
177
158
178
# Main
@@ -188,13 +208,11 @@ def submit_sample(filepath):
188
208
)
189
209
parser .add_argument ("--debug" , action = "store_true" , help = "Enable debug logging." )
190
210
191
- global args
192
211
args = parser .parse_args ()
193
212
194
213
if args .tls :
195
214
schema = "https"
196
215
197
- global api_endpoint
198
216
api_endpoint = "{}://{}:{}/api/checkAsync" .format (schema , args .server , args .port )
199
217
200
218
print ("=" * 80 )
@@ -214,8 +232,8 @@ def submit_sample(filepath):
214
232
print ("Starting the walk at: {} ..." .format (", " .join (args .dirs )))
215
233
216
234
# Walk directory
217
- for dir in args .dirs :
218
- process_dir (dir )
235
+ for walkdir in args .dirs :
236
+ process_dir (walkdir )
219
237
220
238
# End message
221
239
end_date = time .time ()
0 commit comments