1
|
|
|
import os |
2
|
|
|
from skf import settings |
3
|
|
|
from shutil import copyfile |
4
|
|
|
from flask import Flask |
5
|
|
|
from sqlite3 import dbapi2 as sqlite3 |
6
|
|
|
|
7
|
|
|
|
8
|
|
|
app = Flask(__name__) |
9
|
|
|
|
10
|
|
|
def connect_db(): |
11
|
|
|
"""Connects to the specific database.""" |
12
|
|
|
rv = sqlite3.connect(os.path.join(app.root_path, settings.DATABASE)) |
13
|
|
|
rv.row_factory = sqlite3.Row |
14
|
|
|
return rv |
15
|
|
|
|
16
|
|
|
|
17
|
|
|
def init_db(): |
18
|
|
|
"""Initializes the database.""" |
19
|
|
|
try: |
20
|
|
|
os.remove(os.path.join(app.root_path, settings.DATABASE)) |
21
|
|
|
open(os.path.join(app.root_path, 'db.sqlite_schema'), 'a') |
22
|
|
|
os.remove(os.path.join(app.root_path, 'db.sqlite_schema')) |
23
|
|
|
copyfile(os.path.join(app.root_path, "schema.sql"), os.path.join(app.root_path, 'db.sqlite_schema')) |
24
|
|
|
init_md_checklists() |
25
|
|
|
init_md_code_examples() |
26
|
|
|
init_md_knowledge_base() |
27
|
|
|
db = connect_db() |
28
|
|
|
with app.open_resource(os.path.join(app.root_path, 'db.sqlite_schema'), mode='r') as f: |
29
|
|
|
db.cursor().executescript(f.read()) |
30
|
|
|
db.commit() |
31
|
|
|
return True |
32
|
|
|
except: |
33
|
|
|
return False |
34
|
|
|
|
35
|
|
|
|
36
|
|
|
def update_db(): |
37
|
|
|
"""Update the database.""" |
38
|
|
|
try: |
39
|
|
|
os.remove(os.path.join(app.root_path, 'db.sqlite_schema')) |
40
|
|
|
db = connect_db() |
41
|
|
|
db.session.delete("TRUNCATE TABLE kb_items") |
42
|
|
|
db.session.delete("TRUNCATE TABLE code_items") |
43
|
|
|
db.session.delete("TRUNCATE TABLE checklists") |
44
|
|
|
db.session.commit() |
45
|
|
|
|
46
|
|
|
init_md_checklists() |
47
|
|
|
init_md_code_examples() |
48
|
|
|
init_md_knowledge_base() |
49
|
|
|
|
50
|
|
|
with app.open_resource(os.path.join(app.root_path, 'db.sqlite_schema'), mode='r') as f: |
51
|
|
|
db.cursor().executescript(f.read()) |
52
|
|
|
db.commit() |
53
|
|
|
return True |
54
|
|
|
except: |
55
|
|
|
return False |
56
|
|
|
|
57
|
|
|
|
58
|
|
|
def get_db(): |
59
|
|
|
"""Opens a new database connection if there is none yet for the current application context.""" |
60
|
|
|
if not hasattr(g, settings.DATABASE): |
61
|
|
|
g.sqlite_db = connect_db() |
62
|
|
|
return g.sqlite_db |
63
|
|
|
|
64
|
|
|
|
65
|
|
|
def init_md_knowledge_base(): |
66
|
|
|
"""Converts markdown knowledge-base items to DB.""" |
67
|
|
|
kb_dir = os.path.join(app.root_path, 'markdown/knowledge_base') |
68
|
|
|
try: |
69
|
|
|
for filename in os.listdir(kb_dir): |
70
|
|
|
if filename.endswith(".md"): |
71
|
|
|
name_raw = filename.split("-") |
72
|
|
|
kbID = name_raw[0].replace("_", " ") |
73
|
|
|
title = name_raw[3].replace("_", " ") |
74
|
|
|
file = os.path.join(kb_dir, filename) |
75
|
|
|
data = open(file, 'r') |
76
|
|
|
file_content = data.read() |
77
|
|
|
data.close() |
78
|
|
|
content_escaped = file_content.translate(str.maketrans({"'": r"''", "-": r"", "#": r""})) |
79
|
|
|
query = "INSERT OR REPLACE INTO kb_items (kbID, content, title) VALUES ('"+kbID+"','"+content_escaped+"', '"+title+"'); \n" |
80
|
|
|
with open(os.path.join(app.root_path, 'db.sqlite_schema'), 'a') as myfile: |
81
|
|
|
myfile.write(query) |
82
|
|
|
print('Initialized the markdown knowledge-base.') |
83
|
|
|
return True |
84
|
|
|
except: |
85
|
|
|
return False |
86
|
|
|
|
87
|
|
|
|
88
|
|
|
def init_md_code_examples(): |
89
|
|
|
"""Converts markdown code-example items to DB.""" |
90
|
|
|
kb_dir = os.path.join(app.root_path, 'markdown/code_examples/') |
91
|
|
|
code_langs = ['asp', 'java', 'php', 'flask', 'django', 'go', 'ruby'] |
92
|
|
|
try: |
93
|
|
|
for lang in code_langs: |
94
|
|
|
for filename in os.listdir(kb_dir+lang): |
95
|
|
|
if filename.endswith(".md"): |
96
|
|
|
name_raw = filename.split("-") |
97
|
|
|
title = name_raw[3].replace("_", " ") |
98
|
|
|
file = os.path.join(kb_dir+lang, filename) |
99
|
|
|
data = open(file, 'r') |
100
|
|
|
file_content = data.read() |
101
|
|
|
data.close() |
102
|
|
|
content_escaped = file_content.translate(str.maketrans({"'": r"''", "-": r"", "#": r""})) |
103
|
|
|
query = "INSERT OR REPLACE INTO code_items (content, title, code_lang) VALUES ('"+content_escaped+"', '"+title+"', '"+lang+"'); \n" |
104
|
|
|
with open(os.path.join(app.root_path, 'db.sqlite_schema'), 'a') as myfile: |
105
|
|
|
myfile.write(query) |
106
|
|
|
print('Initialized the markdown code-example.') |
107
|
|
|
return True |
108
|
|
|
except: |
109
|
|
|
return False |
110
|
|
|
|
111
|
|
|
|
112
|
|
|
def init_md_checklists(): |
113
|
|
|
"""Converts markdown checklists items to DB.""" |
114
|
|
|
kb_dir = os.path.join(app.root_path, 'markdown/checklists/') |
115
|
|
|
try: |
116
|
|
|
#checklists = ['asvs', 'pcidss', 'custom'] |
117
|
|
|
checklists = ['asvs', 'custom'] |
118
|
|
|
for checklist in checklists: |
119
|
|
|
if checklist == "asvs": |
120
|
|
|
for filename in os.listdir(kb_dir+checklist): |
121
|
|
|
if filename.endswith(".md"): |
122
|
|
|
name_raw = filename.split("-") |
123
|
|
|
level = name_raw[4].replace("_", " ") |
124
|
|
|
kbid_raw = name_raw[6].split(".") |
125
|
|
|
kb_id = kbid_raw[0] |
126
|
|
|
if level == "0": |
127
|
|
|
# For the ASVS categories |
128
|
|
|
file = os.path.join(kb_dir+checklist, filename) |
129
|
|
|
data = open(file, 'r') |
130
|
|
|
file_content = data.read() |
131
|
|
|
data.close() |
132
|
|
|
checklistID_raw = file_content.split(":") |
133
|
|
|
checklistID = checklistID_raw[0] |
134
|
|
|
checklistID = checklistID.lstrip('V') |
135
|
|
|
checklistID = checklistID+".0" |
136
|
|
|
else : |
137
|
|
|
# For the ASVS items |
138
|
|
|
file = os.path.join(kb_dir+checklist, filename) |
139
|
|
|
data = open(file, 'r') |
140
|
|
|
file_content = data.read() |
141
|
|
|
data.close() |
142
|
|
|
checklistID_raw = file_content.split(" ") |
143
|
|
|
checklistID = checklistID_raw[0] |
144
|
|
|
file = os.path.join(kb_dir+checklist, filename) |
145
|
|
|
data = open(file, 'r') |
146
|
|
|
file_content = data.read() |
147
|
|
|
data.close() |
148
|
|
|
content = file_content.split(' ', 1)[1] |
149
|
|
|
content_escaped = content.translate(str.maketrans({"'": r"''", "-": r"", "#": r""})) |
150
|
|
|
query = "INSERT OR REPLACE INTO checklists (checklistID, content, level, kbID) VALUES ('"+checklistID+"', '"+content_escaped+"', '"+level+"', '"+kb_id+"'); \n" |
151
|
|
|
with open(os.path.join(app.root_path, 'db.sqlite_schema'), 'a') as myfile: |
152
|
|
|
myfile.write(query) |
153
|
|
|
print('Initialized the markdown checklists.') |
154
|
|
|
return True |
155
|
|
|
except: |
156
|
|
|
return False |