Skip to content

Commit

Permalink
demoting a few functions to private (mitre#632)
Browse files Browse the repository at this point in the history 
* demoting a few functions to private

* adding comments to sql
  • Loading branch information
david authored Oct 22, 2019
1 parent fe5276c commit e46e203
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 104 deletions.
150 changes: 60 additions & 90 deletions app/service/data_svc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ async def load_adversaries(self, directory):
for pack in [await self._add_adversary_packs(p) for p in adv.get('packs', [])]:
phases += pack
if adv.get('visible', True):
await self.create_adversary(adv['id'], adv['name'], adv['description'], phases)
await self._create_adversary(adv['id'], adv['name'], adv['description'], phases)

async def load_facts(self, directory):
"""
Expand All @@ -69,7 +69,7 @@ async def load_facts(self, directory):

for rule in source.get('rules', []):
rule['source_id'] = source_id
await self.create_rule(**rule)
await self._create_rule(**rule)

async def load_planner(self, directory):
"""
Expand Down Expand Up @@ -120,66 +120,10 @@ async def persist_adversary(self, i, name, description, phases):
p[ability['phase']].append(ability['id'])
f.write(yaml.dump(dict(id=i, name=name, description=description, phases=dict(p))))
f.truncate()
return await self.create_adversary(i, name, description, phases)
return await self._create_adversary(i, name, description, phases)

""" CREATE """

async def create_ability(self, ability_id, tactic, technique_name, technique_id, name, test, description, executor,
platform, cleanup=None, payload=None, parsers=None, requirements=None):
"""
Save a new ability to the database
:param ability_id:
:param tactic:
:param technique_name:
:param technique_id:
:param name:
:param test:
:param description:
:param executor:
:param platform:
:param cleanup:
:param payload:
:param parsers:
:return: the database id
"""
ability = dict(ability_id=ability_id, name=name, test=test, tactic=tactic,
technique_id=technique_id, technique_name=technique_name,
executor=executor, platform=platform, description=description,
cleanup=cleanup)
# update
unique_criteria = dict(ability_id=ability_id, platform=platform, executor=executor)
for entry in await self.dao.get('core_ability', unique_criteria):
await self.update('core_ability', 'id', entry['id'], ability)
for parser in await self.dao.get('core_parser', dict(ability=entry['id'])):
await self.dao.delete('core_parser_map', dict(parser_id=parser['id']))
for requirement in await self.dao.get('core_requirement', dict(ability=entry['id'])):
await self.dao.delete('core_requirement_map', dict(requirement_id=requirement['id']))
await self.dao.delete('core_parser', dict(ability=entry['id']))
await self.dao.delete('core_payload', dict(ability=entry['id']))
return await self._save_ability_extras(entry['id'], payload, parsers, requirements)

# new
identifier = await self.dao.create('core_ability', ability)
return await self._save_ability_extras(identifier, payload, parsers, requirements)

async def create_adversary(self, i, name, description, phases):
"""
Save a new adversary to the database
:param i:
:param name:
:param description:
:param phases:
:return: the database id
"""
identifier = await self.dao.create('core_adversary',
dict(adversary_id=i, name=name, description=description))

await self.dao.delete('core_adversary_map', data=dict(adversary_id=i))
for ability in phases:
a = dict(adversary_id=i, phase=ability['phase'], ability_id=ability['id'])
await self.dao.create('core_adversary_map', a)
return identifier

async def create_operation(self, name, group, adversary_id, jitter='2/8', sources=[],
planner=None, state=None, allow_untrusted=False, autonomous=True):
"""
Expand Down Expand Up @@ -218,21 +162,6 @@ async def create_fact(self, property, value, source_id, score=1, link_id=None):
return await self.dao.create('core_fact', dict(property=property, value=value, source_id=source_id,
score=score, link_id=link_id))

async def create_rule(self, fact, source_id, action='DENY', match='.*'):
"""
Create fact rule. White list or black list. Order matters, executes like firewall rules.
:param source_id: ties to source of rule
:param fact: name of the fact (file.sensitive.extension)
:param action: ALLOW or DENY
:param match: regex or subnet
"""
try:
action = RuleAction[action.upper()].value
await self.dao.create('core_rule', dict(fact=fact, source_id=source_id, action=action, match=match))
except KeyError:
self.log.error(
'Rule action must be in [%s] not %s' % (', '.join(RuleAction.__members__.keys()), action.upper()))

async def create_agent(self, agent, executors):
"""
Save a new agent to the database
Expand Down Expand Up @@ -443,8 +372,8 @@ async def update_agent_executor(self, agent_id, new_executors, previous_executor
:param previous_executors: list of dict with previous executors and their preferred status
:return: None
"""
old_executors = [d['executor'] for d in(sorted(previous_executors, key = lambda i: i['preferred'],
reverse = True))]
old_executors = [d['executor'] for d in (sorted(previous_executors, key=lambda i: i['preferred'],
reverse=True))]

for item in set(new_executors) - set(old_executors):
await self.dao.create('core_executor', dict(agent_id=agent_id, executor=item, preferred=0))
Expand All @@ -458,6 +387,46 @@ async def update_agent_executor(self, agent_id, new_executors, previous_executor

""" PRIVATE """

async def _create_rule(self, fact, source_id, action='DENY', match='.*'):
try:
action = RuleAction[action.upper()].value
await self.dao.create('core_rule', dict(fact=fact, source_id=source_id, action=action, match=match))
except KeyError:
self.log.error(
'Rule action must be in [%s] not %s' % (', '.join(RuleAction.__members__.keys()), action.upper()))

async def _create_ability(self, ability_id, tactic, technique_name, technique_id, name, test, description, executor,
platform, cleanup=None, payload=None, parsers=None, requirements=None):
ability = dict(ability_id=ability_id, name=name, test=test, tactic=tactic,
technique_id=technique_id, technique_name=technique_name,
executor=executor, platform=platform, description=description,
cleanup=cleanup)
# update
unique_criteria = dict(ability_id=ability_id, platform=platform, executor=executor)
for entry in await self.dao.get('core_ability', unique_criteria):
await self.update('core_ability', 'id', entry['id'], ability)
for parser in await self.dao.get('core_parser', dict(ability=entry['id'])):
await self.dao.delete('core_parser_map', dict(parser_id=parser['id']))
for requirement in await self.dao.get('core_requirement', dict(ability=entry['id'])):
await self.dao.delete('core_requirement_map', dict(requirement_id=requirement['id']))
await self.dao.delete('core_parser', dict(ability=entry['id']))
await self.dao.delete('core_payload', dict(ability=entry['id']))
return await self._save_ability_extras(entry['id'], payload, parsers, requirements)

# new
identifier = await self.dao.create('core_ability', ability)
return await self._save_ability_extras(identifier, payload, parsers, requirements)

async def _create_adversary(self, i, name, description, phases):
identifier = await self.dao.create('core_adversary',
dict(adversary_id=i, name=name, description=description))

await self.dao.delete('core_adversary_map', data=dict(adversary_id=i))
for ability in phases:
a = dict(adversary_id=i, phase=ability['phase'], ability_id=ability['id'])
await self.dao.create('core_adversary_map', a)
return identifier

@staticmethod
async def _sort_rules_by_fact(rules):
organized_rules = defaultdict(list)
Expand All @@ -473,18 +442,18 @@ async def _save_ability_to_database(self, filename):
for name, info in executors.items():
for e in name.split(','):
encoded_test = b64encode(info['command'].strip().encode('utf-8'))
await self.create_ability(ability_id=ab.get('id'), tactic=ab['tactic'].lower(),
technique_name=ab['technique']['name'],
technique_id=ab['technique']['attack_id'],
test=encoded_test.decode(),
description=ab.get('description') or '',
executor=e, name=ab['name'], platform=pl,
cleanup=b64encode(
info['cleanup'].strip().encode(
'utf-8')).decode() if info.get(
'cleanup') else None,
payload=info.get('payload'), parsers=info.get('parsers', []),
requirements=ab.get('requirements', []))
await self._create_ability(ability_id=ab.get('id'), tactic=ab['tactic'].lower(),
technique_name=ab['technique']['name'],
technique_id=ab['technique']['attack_id'],
test=encoded_test.decode(),
description=ab.get('description') or '',
executor=e, name=ab['name'], platform=pl,
cleanup=b64encode(
info['cleanup'].strip().encode(
'utf-8')).decode() if info.get(
'cleanup') else None,
payload=info.get('payload'), parsers=info.get('parsers', []),
requirements=ab.get('requirements', []))
await self._delete_stale_abilities(ab)

async def _save_ability_extras(self, identifier, payload, parsers, requirements):
Expand All @@ -500,7 +469,8 @@ async def _save_ability_relationships(self, identifier, table, id_type, relation
for module in relationships:
_id = await self.dao.create(table, dict(ability=identifier, module=module))
for r in relationships.get(module):
relationship = {id_type: _id, 'source': r.get('source'), 'edge': r.get('edge'), 'target': r.get('target')}
relationship = {id_type: _id, 'source': r.get('source'), 'edge': r.get('edge'),
'target': r.get('target')}
await self.dao.create('%s_map' % table, relationship)

async def _delete_stale_abilities(self, ability):
Expand All @@ -519,4 +489,4 @@ async def _add_adversary_packs(self, pack):
async def _add_fact_relationships(self, criteria=None):
relationships = await self.dao.get('core_relationships', criteria)
return [dict(edge=r.get('edge'), target=(await self.dao.get('core_fact', dict(id=r.get('target'))))[0])
for r in relationships if r.get('target')]
for r in relationships if r.get('target')]
1 change: 0 additions & 1 deletion app/service/parsing_svc.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
from base64 import b64decode
from importlib import import_module

from app.service.base_service import BaseService

Expand Down
1 change: 0 additions & 1 deletion app/service/planning_svc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
import re
from base64 import b64decode
from datetime import datetime
from importlib import import_module

from app.service.base_service import BaseService
from app.utility.rule import RuleSet
Expand Down
30 changes: 18 additions & 12 deletions conf/core.sql
View file
Original file line number Diff line number Diff line change
@@ -1,20 +1,26 @@
CREATE TABLE if not exists core_result (link_id integer, output text, parsed data);
/** abilities **/
CREATE TABLE if not exists core_ability (id integer primary key AUTOINCREMENT, ability_id text, tactic text, technique_name, technique_id text, name text, test text, description text, cleanup text, executor, platform, UNIQUE (ability_id, platform, executor) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_payload (ability integer, payload text, UNIQUE (ability, payload) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_adversary (id integer primary key AUTOINCREMENT, adversary_id text, name text, description text, UNIQUE (adversary_id));
CREATE TABLE if not exists core_adversary_map (id integer primary key AUTOINCREMENT, phase integer, adversary_id text, ability_id text, UNIQUE (adversary_id, phase, ability_id));
CREATE TABLE if not exists core_agent (id integer primary key AUTOINCREMENT, paw text, last_seen date, architecture text, platform text, server text, host_group text, location text, pid integer, ppid integer, trusted integer, last_trusted_seen date, sleep_min integer, sleep_max integer);
CREATE TABLE if not exists core_executor (id integer primary key AUTOINCREMENT, agent_id integer, executor text, preferred integer, UNIQUE(agent_id, executor) ON CONFLICT REPLACE);
CREATE TABLE if not exists core_operation (id integer primary key AUTOINCREMENT, name text, host_group text, adversary_id text, jitter text, start date, finish date, phase integer, autonomous integer, planner integer, state text, allow_untrusted integer);
CREATE TABLE if not exists core_chain (id integer primary key AUTOINCREMENT, op_id integer, paw text, ability integer, jitter integer, command text, executor text, cleanup integer, score integer, status integer, decide date, collect date, finish date, UNIQUE(op_id, paw, command));
CREATE TABLE if not exists core_parser (id integer primary key AUTOINCREMENT, ability integer, module text, UNIQUE(ability, module) ON CONFLICT REPLACE);
CREATE TABLE if not exists core_parser_map (parser_id integer, source text, edge text, target text, UNIQUE(parser_id, source, edge, target) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_fact (id integer primary key AUTOINCREMENT, property text, value text, score integer, source_id text, link_id integer DEFAULT 0);
CREATE TABLE if not exists core_used (id integer primary key AUTOINCREMENT, link_id integer, fact_id integer, UNIQUE(link_id, fact_id) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_requirement (id integer primary key AUTOINCREMENT, ability integer, module text, UNIQUE(ability, module) ON CONFLICT REPLACE);
CREATE TABLE if not exists core_requirement_map (requirement_id integer, source text, edge text, target text, UNIQUE(requirement_id, source, edge, target) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_relationships (link_id integer, source integer, edge text, target integer, UNIQUE(link_id, source, edge, target) ON CONFLICT IGNORE);
/** planners **/
CREATE TABLE if not exists core_planner (id integer primary key AUTOINCREMENT, name text, module text, params json, UNIQUE(name) ON CONFLICT IGNORE);
/** agents **/
CREATE TABLE if not exists core_agent (id integer primary key AUTOINCREMENT, paw text, last_seen date, architecture text, platform text, server text, host_group text, location text, pid integer, ppid integer, trusted integer, last_trusted_seen date, sleep_min integer, sleep_max integer);
CREATE TABLE if not exists core_executor (id integer primary key AUTOINCREMENT, agent_id integer, executor text, preferred integer, UNIQUE(agent_id, executor) ON CONFLICT REPLACE);
/** adversaries **/
CREATE TABLE if not exists core_adversary (id integer primary key AUTOINCREMENT, adversary_id text, name text, description text, UNIQUE (adversary_id));
CREATE TABLE if not exists core_adversary_map (id integer primary key AUTOINCREMENT, phase integer, adversary_id text, ability_id text, UNIQUE (adversary_id, phase, ability_id));
/** facts **/
CREATE TABLE if not exists core_fact (id integer primary key AUTOINCREMENT, property text, value text, score integer, source_id text, link_id integer DEFAULT 0);
CREATE TABLE if not exists core_source (id integer primary key AUTOINCREMENT, name text, UNIQUE(name) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_source_map (id integer primary key AUTOINCREMENT, op_id integer, source_id integer, UNIQUE(op_id, source_id) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_planner (id integer primary key AUTOINCREMENT, name text, module text, params json, UNIQUE(name) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_rule (id integer primary key AUTOINCREMENT, action integer, fact text, match text, source_id text);
CREATE TABLE if not exists core_requirement (id integer primary key AUTOINCREMENT, ability integer, module text, UNIQUE(ability, module) ON CONFLICT REPLACE);
CREATE TABLE if not exists core_requirement_map (requirement_id integer, source text, edge text, target text, UNIQUE(requirement_id, source, edge, target) ON CONFLICT IGNORE);
/** operations **/
CREATE TABLE if not exists core_operation (id integer primary key AUTOINCREMENT, name text, host_group text, adversary_id text, jitter text, start date, finish date, phase integer, autonomous integer, planner integer, state text, allow_untrusted integer);
CREATE TABLE if not exists core_chain (id integer primary key AUTOINCREMENT, op_id integer, paw text, ability integer, jitter integer, command text, executor text, cleanup integer, score integer, status integer, decide date, collect date, finish date, UNIQUE(op_id, paw, command));
CREATE TABLE if not exists core_used (id integer primary key AUTOINCREMENT, link_id integer, fact_id integer, UNIQUE(link_id, fact_id) ON CONFLICT IGNORE);
CREATE TABLE if not exists core_result (link_id integer, output text, parsed data);

0 comments on commit e46e203

Please sign in to comment.