-
Notifications
You must be signed in to change notification settings - Fork 0
/
cwetree.py
executable file
·154 lines (121 loc) · 4.6 KB
/
cwetree.py
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
150
151
152
153
154
#!/usr/bin/env python3
# CWEtree v0.3
# https://github.com/mh0x/cwetree
import argparse
import inspect
import io
import json
import os
import requests
import sys
import xml.etree.ElementTree
import zipfile
__version__ = '0.3'
__author__ = 'https://github.com/mh0x'
script_name = 'CWEtree v' + __version__ + ' (' + __author__ + '/cwetree)'
cwe_xml_url = 'http://cwe.mitre.org/cwe-6'
cwe_views_url = 'https://cwe.mitre.org/data/xml/views/'
cwe_views = {'699': 'Development Concepts',
'1000': 'Research Concepts',
'1008': 'Architectural Concepts'}
script_path = os.path.realpath(inspect.getsourcefile(lambda: 0))
default_dir_rel = os.path.join('html', 'js')
default_dir = os.path.join(os.path.dirname(script_path), default_dir_rel)
def error(err):
print('[!] error: ' + str(err), file=sys.stderr)
sys.exit(1)
def info(msg, quiet=False):
if not quiet:
print('[*] ' + msg)
def prologue(quiet=False):
if not quiet:
print(script_name)
class ArgParser(argparse.ArgumentParser):
def format_help(self):
formatter = self._get_formatter()
formatter.add_text(script_name)
formatter.add_usage(self.usage, self._actions,
self._mutually_exclusive_groups)
for action_group in self._action_groups:
formatter.start_section(action_group.title)
formatter.add_text(action_group.description)
formatter.add_arguments(action_group._group_actions)
formatter.end_section()
formatter.add_text(self.epilog)
return formatter.format_help()
def error(self, msg):
raise argparse.ArgumentTypeError(msg + os.linesep*2
+ self.format_usage().rstrip())
def parse_args():
parser = ArgParser()
parser.add_argument('-q', '--quiet', action='store_true',
help='suppress messages sent to stdout')
parser.add_argument('-d', '--dir', default=default_dir,
help='output directory path (default: '
+ default_dir_rel + ')')
try:
return parser.parse_args()
except argparse.ArgumentTypeError as err:
prologue()
print()
error(err)
def download_xml(view_id, quiet=False):
xfile = view_id + '.xml'
zfile = xfile + '.zip'
info('downloading ' + zfile, quiet)
try:
resp = requests.get(cwe_views_url + zfile)
info('unzipping ' + zfile, quiet)
with zipfile.ZipFile(io.BytesIO(resp.content), 'r') as zip:
with zip.open(xfile, 'r') as fp:
info('parsing ' + xfile, quiet)
return xml.etree.ElementTree.parse(fp).getroot()
except (IOError, zipfile.BadZipFile) as err:
error(err)
def generate_node(id, name, children=[]):
if children:
return {'cid': id, 'name': name, 'children': children}
return {'cid': id, 'name': name}
def generate_children(id, items):
children = []
for id1 in items:
if id in items[id1]['parents']:
children.append(generate_node(id1, items[id1]['name'],
generate_children(id1, items)))
return children
def generate_tree(view_id, view_name, items):
tree = generate_node(view_id, view_name)
tree['children'] = []
for id in [id for id in items if not items[id]['parents']]:
tree['children'].append(generate_node(id, items[id]['name'],
generate_children(id, items)))
return tree
def generate_json(view_id, view_name, view_xml, dir=default_dir, quiet=False):
jfile = view_id + '.js'
info('generating ' + jfile, quiet)
items = {}
for item in view_xml.findall('.//{' + cwe_xml_url + '}Weakness'):
id = item.get('ID')
related = item.find('{' + cwe_xml_url + '}Related_Weaknesses')
items[id] = {'id': id, 'name': item.get('Name'), 'parents': {
item.get('CWE_ID') for item in related
if item.get('Nature') == 'ChildOf'} if related else set()}
try:
with open(os.path.join(dir, jfile), 'w') as fp:
fp.write('var view_' + view_id + ' = ')
json.dump(generate_tree(view_id, view_name, items), fp)
except IOError as err:
error(err)
def main():
args = parse_args()
prologue(args.quiet)
try:
for view_id, view_name in cwe_views.items():
view_xml = download_xml(view_id, args.quiet)
generate_json(view_id, view_name, view_xml, args.dir, args.quiet)
except KeyboardInterrupt:
print(' stopping ...')
else:
info('done', args.quiet)
if __name__ == '__main__':
main()