import time
import json
import pydgraph
from pydgraph import DgraphClient, DgraphClientStub
def set_process_schema(client: DgraphClient, engagement: bool = False):
schema = """node_key: string @upsert @index(hash) .
process_id: int @index(int) .
created_timestamp: int @index(int) .
asset_id: string @index(hash) .
terminate_time: int @index(int) .
image_name: string @index(exact, hash, trigram, fulltext) .
process_name: string @index(exact, hash, trigram, fulltext) .
arguments: string @index(fulltext) @index(trigram) .
bin_file: uid @reverse .
children: [uid] @reverse .
created_files: [uid] @reverse .
deleted_files: [uid] @reverse .
read_files: [uid] @reverse .
wrote_files: [uid] @reverse .
created_connections: [uid] @reverse .
bound_connections: [uid] @reverse .
"""
if engagement:
schema += "\n"
schema += "risks: [uid] @reverse ."
schema += """
process_guid: string @index(exact, hash, trigram, fulltext) .
"""
op = pydgraph.Operation(schema=schema)
client.alter(op)
def set_file_schema(client: DgraphClient, engagement: bool=False) -> None:
schema = """
node_key: string @upsert @index(hash) .
file_name: string @index(exact, hash, trigram, fulltext) .
asset_id: string @index(exact, hash, trigram, fulltext) .
file_path: string @index(exact, hash, trigram, fulltext) .
file_extension: string @index(exact, hash, trigram, fulltext) .
file_mime_type: string @index(exact, hash, trigram, fulltext) .
file_size: int @index(int) .
file_version: string @index(exact, hash, trigram, fulltext) .
file_description: string @index(exact, hash, trigram, fulltext) .
file_product: string @index(exact, hash, trigram, fulltext) .
file_company: string @index(exact, hash, trigram, fulltext) .
file_directory: string @index(exact, hash, trigram, fulltext) .
file_inode: int @index(int) .
file_hard_links: string @index(exact, hash, trigram, fulltext) .
signed: bool @index(bool) .
signed_status: string @index(exact, hash, trigram, fulltext) .
md5_hash: string @index(exact, hash, trigram, fulltext) .
sha1_hash: string @index(exact, hash, trigram, fulltext) .
sha256_hash: string @index(exact, hash, trigram, fulltext) .
"""
if engagement:
schema += "\n"
schema += "risks: uid @reverse ."
op = pydgraph.Operation(schema=schema)
client.alter(op)
def set_outbound_connection_schema(client, engagement=False):
schema = """
create_time: int @index(int) .
terminate_time: int @index(int) .
last_seen_time: int @index(int) .
ip: string @index(exact, trigram, hash) .
port: string @index(exact, trigram, hash) .
"""
if engagement:
schema += "\n"
schema += "risks: [uid] @reverse ."
op = pydgraph.Operation(schema=schema)
client.alter(op)
def set_inbound_connection_schema(client, engagement=False):
schema = """
node_key: string @upsert @index(hash) .
asset_id: string @index(exact, hash, trigram, fulltext) .
port: string @index(exact, trigram, hash) .
"""
if engagement:
schema += "\n"
schema += "risks: [uid] @reverse ."
op = pydgraph.Operation(schema=schema)
client.alter(op)
def set_external_ip_schema(client, engagement=False):
schema = """
node_key: string @upsert @index(hash) .
external_ip: string @index(exact, trigram, hash) .
"""
if engagement:
schema += "\n"
schema += "risks: [uid] @reverse ."
op = pydgraph.Operation(schema=schema)
client.alter(op)
def set_risk_schema(client, engagement=False):
schema = """
analyzer_name: string @index(exact, trigram, hash) .
risk_score: int @index(int) .
"""
op = pydgraph.Operation(schema=schema)
client.alter(op)
def set_lens_schema(client, engagement=False):
schema = """
scope: [uid] @reverse .
lens: string @upsert @index(exact, trigram, hash) .
score: int @index(int) .
"""
op = pydgraph.Operation(schema=schema)
client.alter(op)
def set_ipc_schema(client, engagement=False):
schema = """
node_key: string @upsert @index(hash) .
node_type: string @index(hash) .
key: string @index(hash) .
node_type: string @index(hash) .
ipc_type: string @index(hash) .
src_pid: int @index(int) .
dst_pid: int @index(int) .
ipc_creator: uid @reverse .
ipc_recipient: uid @reverse .
"""
if engagement:
schema += "\n"
schema += "risks: [uid] @reverse ."
op = pydgraph.Operation(schema=schema)
client.alter(op)
def drop_all(client):
op = pydgraph.Operation(drop_all=True)
client.alter(op)
client = DgraphClient(DgraphClientStub('localhost:9080'))
drop_all(client)
set_process_schema(client)
set_file_schema(client)
set_outbound_connection_schema(client)
set_inbound_connection_schema(client)
set_external_ip_schema(client)
set_ipc_schema(client)