SalSA/main.py at master · deptofdefense/SalSA · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
"""
Generates report from alerts generated by rules
"""
import os
import sys
import pprint
import logging
import hashlib
import datetime
import argparse

# import PE parser class
import pe

# import all rules
from rules import check_imported_dlls
from rules import check_section_permissions
from rules import check_section_names
from rules import check_section_sizes
from rules import check_section_entropy
from rules import check_section_strings
from rules import check_dos_stub


if __name__ == '__main__':
  # parse user arguments
  parser = argparse.ArgumentParser()
  parser.add_argument('-f',      dest='file',      help='file to parse',                   type=str)
  parser.add_argument('-d',      dest='directory', help='directory of files to parse',     type=str)
  parser.add_argument('-l',      dest='log_level', help='logging level (default is INFO)', type=str, default='INFO', choices=['CRITICAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG'])
  parser.add_argument('--html',  dest='html',      help='generate HTML output',            action='store_true')
  parser.add_argument('--debug', dest='debug',     help='display full PE dump',            action='store_true')
  args = parser.parse_args()
  logging.basicConfig(format='[%(asctime)s][%(levelname)s] %(message)s', level=args.log_level, datefmt='%d-%b-%y %H:%M:%S')
  # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  # parse each file specified
  # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  def _hash(filename, block_size=1024):
    """ used to get the sha256 hash of a file for reporting """
    sha256 = hashlib.sha256()
    sha1 = hashlib.sha1()
    md5 = hashlib.md5()
    with open(filename, 'rb') as f:
      for b in iter(lambda: f.read(block_size), b''):
        sha256.update(b)
        sha1.update(b)
        md5.update(b)
    return {
      'md5': md5.hexdigest(),
      'sha1': sha1.hexdigest(),
      'sha256': sha256.hexdigest()
    }
  pedata = []
  if args.file and os.path.isfile(args.file):
    logging.info('parsing file {0} ...'.format(args.file))
    pedata.append({
      'file': os.path.basename(args.file),
      'data': pe.PE(args.file),
      'hashes': _hash(args.file),
    })
  elif args.directory:
    for f in os.listdir(args.directory):
      path = os.path.join(args.directory, f)
      if os.path.isfile(path):
        logging.info('parsing file {0}'.format(path))
        pedata.append({
          'file': f,
          'data': pe.PE(path),
          'hashes': _hash(path),
        })
  else:
    logging.error('no files specified for parsing.')
    quit()
  # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  # apply rules
  # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  def _run(r, d):
    """ used to run a specific rule within its own error context """
    try:
      return sys.modules['rules.' + r].run(d)
    except Exception as e:
      logging.error('Failed to run rule "{0}" due to error: {1}'.format(r, str(e)))
      return []
  for obj in pedata:
    alerts = []
    logging.info('applying rules to {0} ...'.format(obj['file']))
    alerts.extend(_run('check_imported_dlls', obj['data']))
    alerts.extend(_run('check_section_permissions', obj['data']))
    alerts.extend(_run('check_section_names', obj['data']))
    alerts.extend(_run('check_section_sizes', obj['data']))
    alerts.extend(_run('check_section_entropy', obj['data']))
    alerts.extend(_run('check_section_strings', obj['data']))
    alerts.extend(_run('check_dos_stub', obj['data']))
    obj['alerts'] = alerts
  # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  # generate output
  # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  if args.html:
    # HTML output template
    HTML_TEMPLATE = """
      <!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><title>{title}</title></head>
      <body style="font-family: monospace; font-size: 1em; color: #17202a">
      <div style="margin: 10px auto;">{alerts}</div></body></html>"""
    filename = 'salsa_report_{0}.html'.format(datetime.datetime.now().strftime('%b-%d-%Y-%Hh-%Mm-%Ss'))
    logging.info('generating report {0} ...'.format(filename))
    # parse each file's results
    html = ''
    for f in pedata:
      html += '<h2 style="color: #c0392b;">{0}</h2>'.format(f['file'])
      html += '<p>sha256: {0}</p>'.format(f['hashes']['sha256'])
      html += '<p>sha1: {0}</p>'.format(f['hashes']['sha1'])
      html += '<p>md5: {0}</p>'.format(f['hashes']['md5'])
      html += '<hr>'
      # parse alerts for this file
      for d in f['alerts']:
        html += '<div>'
        html += '<h4 style="color: #d35400;">{0}</h4>'.format(d['title'])
        html += '<p>{0}</p>'.format(d['description'])
        if d['data']:
          html += '<ul style="list-style-type: square; color: #566573;">'
          for i in d['data']:
            html += '<li>{0}</li>'.format(i)
          html += '</ul>'
        if d['code']:
          html += '<pre style="background-color: #eaecee; color: #3498db">{0}</pre>'.format(d['code'])
        html += '</div>'
    with open(filename, 'w') as f:
      f.write(HTML_TEMPLATE.format(title=filename, alerts=html))
  else:
    # display alerts
    for f in pedata:
      print('[{0}] SHA256 {1}'.format(f['file'], f['hashes']['sha256']))
      print('[{0}] SHA1 {1}'.format(f['file'], f['hashes']['sha1']))
      print('[{0}] MD5 {1}'.format(f['file'], f['hashes']['md5']))
      print('-' * 20) # visual delimiter
      for d in f['alerts']:
        print('[{0}] Alert: {1}'.format(f['file'], d['title']))
        print('[{0}] Description: {1}'.format(f['file'], d['description']))
        if d['data']:
          for i in d['data']:
            print('[{0}] * {1}'.format(f['file'], i))
        if d['code']:
          print('[{0}] Code:'.format(f['file']))
          print(d['code'])
        print('-' * 20) # visual delimiter
      # check for debugging flag
      if args.debug:
        print('[{0}] Full PE Dump:'.format(f['file']))
        pprint.pprint(f['data'].hex(f['data'].dict()), indent=2)