Peter M. Groen / oletools

Commit c203e02dc9bab6fb7affc6012d80e27a024abdd3

Authored by Philippe Lagadec
1 parent efe28a7b

mraptor: now returns an exit code based on the overall result

Inline Side-by-side
Showing 1 changed file with 63 additions and 26 deletions
oletools/mraptor.py
  View file @c203e02
... ... @@ -48,8 +48,9 @@ http://www.decalage.info/python/oletools
48 48 # CHANGELOG:
49 49 # 2016-02-23 v0.01 PL: - first version
50 50 # 2016-02-29 v0.02 PL: - added Workbook_Activate, FileSaveAs
  51 +# 2016-03-04 v0.03 PL: - returns an exit code based on the overall result
51 52  
52   -__version__ = '0.01'
  53 +__version__ = '0.03'
53 54  
54 55 #------------------------------------------------------------------------------
55 56 # TODO:
... ... @@ -113,6 +114,36 @@ TYPE2TAG = {
113 114  
114 115 # === CLASSES =================================================================
115 116  
  117 +class Result_NoMacro(object):
  118 + exit_code = 0
  119 + color = 'green'
  120 + name = 'No Macro'
  121 +
  122 +
  123 +class Result_NotMSOffice(object):
  124 + exit_code = 1
  125 + color = 'green'
  126 + name = 'Not MS Office'
  127 +
  128 +
  129 +class Result_MacroOK(object):
  130 + exit_code = 2
  131 + color = 'cyan'
  132 + name = 'Macro OK'
  133 +
  134 +
  135 +class Result_Error(object):
  136 + exit_code = 10
  137 + color = 'yellow'
  138 + name = 'ERROR'
  139 +
  140 +
  141 +class Result_Suspicious(object):
  142 + exit_code = 20
  143 + color = 'red'
  144 + name = 'SUSPICIOUS'
  145 +
  146 +
116 147 class MacroRaptor(object):
117 148 """
118 149 class to scan VBA macro code to detect if it is malicious
... ... @@ -202,6 +233,9 @@ def main():
202 233 if len(args) == 0:
203 234 print __doc__
204 235 parser.print_help()
  236 + print '\nAn exit code is returned based on the analysis result:'
  237 + for result in (Result_NoMacro, Result_NotMSOffice, Result_MacroOK, Result_Error, Result_Suspicious):
  238 + print ' - %d: %s' % (result.exit_code, result.name)
205 239 sys.exit()
206 240  
207 241 # print banner with version
... ... @@ -216,6 +250,8 @@ def main():
216 250 header_row=['Result', 'Flags', 'Type', 'File'],
217 251 column_width=[10, 5, 4, 56])
218 252  
  253 + exitcode = -1
  254 + global_result = None
219 255 # TODO: handle errors in xglob, to continue processing the next files
220 256 for container, filename, data in xglob.iter_files(args, recursive=options.recursive,
221 257 zip_password=options.zip_password, zip_fname=options.zip_fname):
... ... @@ -231,24 +267,24 @@ def main():
231 267 # log.exception('Error when opening file %r' % full_name)
232 268 # continue
233 269 if isinstance(data, Exception):
234   - result = '* ERROR *'
235   - result_color = 'yellow'
236   - t.write_row([result, '', '', full_name],
237   - colors=[result_color, None, None, None])
  270 + result = Result_Error
  271 + t.write_row([result.name, '', '', full_name],
  272 + colors=[result.color, None, None, None])
238 273 t.write_row(['', '', '', str(data)],
239   - colors=[None, None, None, result_color])
  274 + colors=[None, None, None, result.color])
240 275 else:
  276 + filetype = '???'
241 277 try:
242 278 vba_parser = olevba.VBA_Parser(filename=filename, data=data, container=container)
243 279 filetype = TYPE2TAG[vba_parser.type]
244 280 except Exception as e:
245 281 # log.error('Error when parsing VBA macros from file %r' % full_name)
246   - result = '* ERROR *'
247   - result_color = 'yellow'
248   - t.write_row([result, '', TYPE2TAG[vba_parser.type], full_name],
249   - colors=[result_color, None, None, None])
  282 + # TODO: distinguish actual errors from non-MSOffice files
  283 + result = Result_Error
  284 + t.write_row([result.name, '', filetype, full_name],
  285 + colors=[result.color, None, None, None])
250 286 t.write_row(['', '', '', str(e)],
251   - colors=[None, None, None, result_color])
  287 + colors=[None, None, None, result.color])
252 288 continue
253 289 if vba_parser.detect_vba_macros():
254 290 vba_code_all_modules = ''
... ... @@ -257,33 +293,34 @@ def main():
257 293 vba_code_all_modules += vba_code + '\n'
258 294 except Exception as e:
259 295 # log.error('Error when parsing VBA macros from file %r' % full_name)
260   - result = '* ERROR *'
261   - result_color = 'yellow'
262   - t.write_row([result, '', TYPE2TAG[vba_parser.type], full_name],
263   - colors=[result_color, None, None, None])
  296 + result = Result_Error
  297 + t.write_row([result.name, '', TYPE2TAG[vba_parser.type], full_name],
  298 + colors=[result.color, None, None, None])
264 299 t.write_row(['', '', '', str(e)],
265   - colors=[None, None, None, result_color])
  300 + colors=[None, None, None, result.color])
266 301 continue
267 302 mraptor = MacroRaptor(vba_code_all_modules)
268 303 mraptor.scan()
269 304 if mraptor.suspicious:
270   - result = 'SUSPICIOUS'
271   - result_color = 'red'
  305 + result = Result_Suspicious
272 306 else:
273   - result = 'Macro OK'
274   - result_color = 'cyan'
275   - t.write_row([result, mraptor.get_flags(), filetype, full_name],
276   - colors=[result_color, None, None, None])
  307 + result = Result_MacroOK
  308 + t.write_row([result.name, mraptor.get_flags(), filetype, full_name],
  309 + colors=[result.color, None, None, None])
277 310 if mraptor.matches and options.show_matches:
278 311 t.write_row(['', '', '', 'Matches: %r' % mraptor.matches])
279 312 else:
280   - result = 'No Macro'
281   - result_color = 'green'
282   - t.write_row([result, '', filetype, full_name],
283   - colors=[result_color, None, None, None])
  313 + result = Result_NoMacro
  314 + t.write_row([result.name, '', filetype, full_name],
  315 + colors=[result.color, None, None, None])
  316 + if result.exit_code > exitcode:
  317 + global_result = result
  318 + exitcode = result.exit_code
284 319  
285 320 print ''
286 321 print 'Flags: A=AutoExec, W=Write, X=Execute'
  322 + print 'Exit code: %d - %s' % (exitcode, global_result.name)
  323 + sys.exit(exitcode)
287 324  
288 325 if __name__ == '__main__':
289 326 main()
... ...