DISARMframeworks/CODE/generate_disarm_sql.py at main · kogratte/DISARMframeworks · 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
''' Generate DISARM sql databases

# Author: SJ Terp, Twitter @bodaceacat
# License: GPL-3

'''

import pandas as pd
import sqlite3 as sql
import os
from sqlalchemy import create_engine
from generate_DISARM_pages import Disarm


def generate_disarm_sql(dbasetype='sqlite'):
	''' Generate SQL
	Expected database types include
		'sqlite': used to export database to other apps
		'postgresql_local': used in local instances of the DISARM explorer
		'heroku': used in Heroku-based instances of the DISARM explorer, although current
		practice at DISARM is to post a dump of the postgresql database to it instead.
	'''

	# Generate DISARM datasets
	disarm = Disarm()


	# Create connection to database
	if dbasetype == 'sqlite':
	    # Generate minimal sqlite database from the  variables
	    conn = sql.connect('../generated_files/DISARM_database.sqlite')
	elif dbasetype == 'postgresql_local':
	    conn = create_engine('postgresql://sara:@localhost:5432/disarmsite')


	# Add table to database
	def add_table(dataframe, tablename, columns):
	    # Create sql-appropriate table from dataframe
	    colnames = ', '.join(['{} TEXT NOT NULL'.format(col) for col in columns])
	    newtable = dataframe[columns].copy().applymap(str)
	    newtable['id'] = range(1,len(newtable)+1)

	    # send to database
	    if dbasetype == 'sqlite':
	        conn.execute("DROP TABLE IF EXISTS {}".format(tablename))
	        conn.execute('''CREATE TABLE {} (id INTEGER PRIMARY KEY AUTOINCREMENT, {});'''.format(tablename, colnames))
	        newtable.to_sql(tablename, conn, index=False, if_exists='append')
	        conn.commit()
	    elif dbasetype == 'postgresql_local':
	        newtable.to_sql(tablename, conn, index=False, if_exists='replace')
	    return newtable


	# Build a cross-table
	def object_tactics_techniques(objectcol, objecttable, crosstable):
	    # objects to techniques
	    ctech = crosstable.copy()
	    ctech = ctech[(ctech['technique_id'] != '') & (~ctech['technique_id'].str.startswith('TA'))]
	    ctech.rename(columns={'disarm_id':objectcol}, inplace=True)
	    ctech['summary'] = 'N/A'

	    # objects to tactics
	    ctact = crosstable[crosstable['technique_id'].str.startswith('TA')].copy()
	    ctact.rename(columns={'disarm_id':objectcol, 'technique_id': 'tactic_id'}, inplace=True)
	    ctact['main_tactic'] = 'N'
	    ctactmain = objecttable[['disarm_id', 'tactic_id']].copy()
	    ctactmain.rename(columns={'disarm_id':objectcol}, inplace=True)
	    ctactmain['main_tactic'] = 'Y'
	    ctact = pd.concat([ctact, ctactmain], ignore_index=True, sort=False)
	    ctact['summary'] = 'N/A'
	    return(ctech, ctact)

	#Load all the tables - Heroku needs them in correct order...

	# -- frameworks --

	#newtable = add_table(disarm.df_actortypes, 'actor_type', ['disarm_id', 'sector_id', 'framework_id', 'name', 'summary'])
	newtable = add_table(disarm.df_counters, 'counter', ['disarm_id', 'tactic_id', 'metatechnique_id', 'name', 'summary'])
	newtable = add_table(disarm.df_detections, 'detection', ['disarm_id', 'tactic_id', 'name', 'summary'])
	newtable = add_table(disarm.df_frameworks, 'framework', ['disarm_id', 'name', 'summary'])
	newtable = add_table(disarm.df_metatechniques, 'metatechnique', ['disarm_id', 'name', 'summary'])
	newtable = add_table(disarm.df_phases, 'phase', ['disarm_id', 'name', 'rank', 'summary'])
	newtable = add_table(disarm.df_playbooks, 'playbook', ['disarm_id', 'object_id', 'name', 'summary'])
	newtable = add_table(disarm.df_resources, 'resource', ['disarm_id', 'name', 'summary', 'resource_type'])
	newtable = add_table(disarm.df_responsetypes, 'responsetype', ['disarm_id', 'name', 'summary'])
	#newtable = add_table(disarm.df_sector, 'sector', ['disarm_id', 'name', 'summary'])
	newtable = add_table(disarm.df_tactics, 'tactic', ['disarm_id', 'phase_id', 'name', 'rank', 'summary'])
	newtable = add_table(disarm.df_tasks, 'task', ['disarm_id', 'tactic_id', 'framework_id', 'name', 'summary'])
	newtable = add_table(disarm.df_techniques, 'technique', ['disarm_id', 'tactic_id', 'name', 'summary'])

	(ctech, ctact) = object_tactics_techniques('counter_id', disarm.df_counters, disarm.cross_counterid_techniqueid)
	newtable = add_table(ctech, 'counter_technique', ['counter_id', 'technique_id', 'summary'])
	newtable = add_table(ctact, 'counter_tactic', ['counter_id', 'tactic_id', 'main_tactic', 'summary'])

	(dtech, dtact) = object_tactics_techniques('detection_id', disarm.df_detections, disarm.cross_detectionid_techniqueid)
	newtable = add_table(dtech, 'detection_technique', ['detection_id', 'technique_id', 'summary'])
	newtable = add_table(dtact, 'detection_tactic', ['detection_id', 'tactic_id', 'main_tactic', 'summary'])

	# -- datasets --

	newtable = add_table(disarm.df_examples, 'example', ['disarm_id', 'object_id', 'name', 'summary'])

	# dataset
	# reference - create this from other tables.
	# Also need incidentcounter etc - create from data
	newtable = add_table(disarm.df_externalgroups, 'externalgroup', ['disarm_id', 'name', 'url', 'summary',
	                                                'sector', 'primary_role', 'secondary_role',
	                                                'primary_subject', 'secondary_subject',
	                                                'volunteers', 'region', 'country',
	                                                'twitter_handle'])
	newtable = add_table(disarm.df_incidents, 'incident', ['disarm_id', 'name', 'summary',
	                                                      'year_started', 'attributions_seen',
	                                                      'found_in_country', 'objecttype'])
	newtable = add_table(disarm.df_tools, 'tool', ['disarm_id', 'name', 'summary',
	                                              'externalgroup', 'url', 'category',
	                                              'disinformation_use', 'cogseccollab_use',
	                                              'function', 'code_url', 'artifacts',
	                                              'automation', 'platform', 'accessibility'])
	# incidenttechnique crosstable
	it = vars(disarm)['it'][['disarm_id', 'name', 'summary', 'disarm_id_incident', 'disarm_id_technique']].copy()
	it.rename(columns={'disarm_id_incident':'incident_id', 'disarm_id_technique': 'technique_id'}, inplace=True)
	newtable = add_table(it, 'incident_technique', ['disarm_id', 'name', 'summary',
	                                                'incident_id', 'technique_id'])

	# Load in users table, and close connection
	if dbasetype == 'sqlite':
	    conn.execute("DROP TABLE IF EXISTS {}".format('users'))
	    conn.execute('''CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT NOT NULL UNIQUE, password TEXT NOT NULL);''')
	    conn.close()
	else:
	    dfusers = pd.DataFrame([['test','testing']], columns=['username', 'password'])
	    usertable = add_table(dfusers, 'users', ['username', 'password'])

	    return


''' main, if we need it
'''
def main():
    generate_disarm_sql()


if __name__ == "__main__":
    main()