def set_schema(schema=None):
if schema is None:
schema = open('./schema.dgraph').read()
print("schema: \n", schema)
op = pydgraph.Operation(schema=schema)
print("set_schema: ", op)
return client.alter(op)
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 ."
# unstable
schema += """
process_guid: string @index(exact, hash, trigram, fulltext) .
"""
op = pydgraph.Operation(schema=schema)
client.alter(op)
def set_schema(client):
schema = """
<EMPLOYED_BY>: [uid] @reverse .
<Name>: string @index(exact) .
<dgraph.graphql.schema>: string .
"""
return client.alter(pydgraph.Operation(schema=schema))
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_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 drop_all(self):
"""
Drops all data from dgraph
:return:
"""
op = pydgraph.Operation(drop_all=True)
self.client.alter(op)
def initGraphTypes(client):
schema = '''
created_date: datetime .
last_modified_date: datetime .
likelihood: string .
prerequisites: [string] .
resources_required: [string] .
severity: string .
attack_technique: [uid] @reverse .
type Capec {
name
created_date
last_modified_date
name_full
description
reference
likelihood
prerequisites
resources_required
severity
attack_technique
}
'''
op = pydgraph.Operation(
schema=schema, run_in_background=False)
client.alter(op)
print("schema initialized")
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 setup_schema(self):
schema = """
email_address: string @index(exact) .
id: string @index(exact) .
"""
op = pydgraph.Operation(schema=schema)
self.dg.alter(op)
def add_xid_to_schema(self):
"""
adds the xid schema to dgraph
:return: True if added
"""
schema = "xid: string @index(exact) ."
op = pydgraph.Operation(schema=schema)
self.client.alter(op)
def set_schema(client):
schema = """
name: string @index(exact) .
age: int .
married: bool .
loc: geo .
dob: datetime .
"""
return client.alter(pydgraph.Operation(schema=schema))
def create_process_schema(eg_client: DgraphClient):
schema = \
'node_key: string @index(hash) .\n' +\
'engagement_key: string @index(hash) .\n' +\
'children: uid @reverse .\n' +\
'pid: int @index(int) .\n'
op = pydgraph.Operation(schema=schema)
eg_client.alter(op)
def set_schema(client, schema=None):
if schema is None:
schema = """
email: string @index(exact) .
name: string @index(term) .
age: int .
referred: uid @count @reverse .
"""
op = pydgraph.Operation(schema=schema)
return client.alter(op)
def rebuild(self):
'''
重新建立表结构
:return:
'''
schema = '''xid: string @index(exact) .
type.user: string @index(exact) .
type.work: string @index(exact) .
friend: uid @count @reverse .
work: uid @count @reverse .
work.created_at: dateTime @index(hour) .
work.status: int @index(int) .
news: uid @count @reverse .
news.type: string @index(exact) .
news.created_at: dateTime @index(hour) .'''
self.client.alter(pydgraph.Operation(drop_all=True))
op = pydgraph.Operation(schema=schema)
self.client.alter(op)
def set_schema(client):
"""
Sets up schema for data we are about to ingest.
:param client:
:return:
"""
print("Initializing DGraph Schema")
return client.alter(pydgraph.Operation(schema=schema))
def main():
client = create_client('localhost:9080')
client.login("groot", "password")
# Drop all
client.alter(pydgraph.Operation(drop_all=True))
# Update schema
schema = '''
name: string @index(exact) .
description: string .
url: string .
'''
op = pydgraph.Operation(schema=schema)
client.alter(op)
# Mutate
dgraph = {
"name": "Dgraph",
"description": "Scalable, Distributed, Low Latency Graph Database",
"url": "https://dgraph.io"
}
txn = client.txn()
try:
txn.mutate(set_obj=dgraph)
txn.commit()
finally:
txn.discard()
# Query
res = client.txn(read_only=True).query('''
query dgraph($name: string) {
data(func: eq(name, $name)) {
uid
name
description
url
}
}
''',
variables={"$name": "Dgraph"})
print(res.json)
def drop_all(client):
"""
Wipe the database for a fresh start
:param client: the dgraph client
:return: dgraph response
"""
try:
LOG.info("Dropping existing graph")
return client.alter(pydgraph.Operation(drop_all=True))
except:
LOG.critical("Failed to drop previous graph")
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_schema(self, schema=None):
# for testing
if not schema:
schema = """
name: string @index(exact) .
friend: uid @reverse .
age: int .
married: bool .
loc: geo .
dob: datetime .
"""
self.client.alter(pydgraph.Operation(schema=schema))
def drop_all(client):
"""
This function drops all dgraph nodes, edges and schemas.
Essentially a reset on DGraph
:param client:
:return:
"""
print("Dropping DGraph data")
return client.alter(pydgraph.Operation(drop_all=True))
def initGraphTypes(client):
schema = '''
cwe: [uid] @reverse .
reference: [uid] .
impact: uid .
published_date: dateTime .
last_modified_date: dateTime .
type Cve {
name
cwe
reference
description
impact
published_date
last_modified_date
}
url: string .
refsource: string .
tag: [string] .
type Reference {
url
name
description
refsource
tag
}
severity: string .
exploitability_score: float .
impact_score: float .
obtain_all_privilege: bool .
obtain_user_privilege: bool .
obtain_other_privilege: bool .
user_interaction_required: bool .
type Impact {
severity
exploitability_score
impact_score
obtain_all_privilege
obtain_user_privilege
obtain_other_privilege
user_interaction_required
}
'''
op = pydgraph.Operation(schema=schema, run_in_background=False)
client.alter(op)
print("schema initialized")
def initGraphTypes(client):
schema = '''
sid: string @index(exact) .
cve: [uid] @reverse .
type SuricataRule {
sid
cve
}
'''
op = pydgraph.Operation(
schema=schema, run_in_background=False)
client.alter(op)
print("schema initialized")
def set_schema(client):
schema = """
firstname: string @index(term) .
lastname: string @index(term) .
age: int @index(int) .
phonenumber: string @index(term) .
type Person {
firstname
lastname
age
phonenumber
}
"""
return client.alter(pydgraph.Operation(schema=schema))
def _schema(self, schema):
#
# Example:
# schema = \
# """
# identity: int @index(int) .
# category: int .
# document_tokens: string .
# cite: uid @count .
# """
#
print(schema)
op = pydgraph.Operation(schema=schema)
self._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_schema(client):
schema = """
name: string @index(exact) .
depends: [uid] @reverse .
version: string .
src: string .
pkg_rel_date: string .
number_dependents: int .
type Package {
name
depends
version
}
"""
return client.alter(pydgraph.Operation(schema=schema))
def add_genome_to_schema(client, genome):
"""
Index the genome name as a predicate, so functions can be used on it when searching etc.
If the genome name is not added to the schema, it will not be indexed, and functions can
not be applied to it. We will use the upsert directive when modifying the schema, so that
if it exists, nothing happens, but new names are added to the schema.
:param client: dgraph client
:param genome: genomeName
:return: Altered dgraph
"""
schema = """
{0}: uid .
""".format(genome)
return client.alter(pydgraph.Operation(schema=schema))
def add_schema(client):
"""
We can define whatever type we like for our schema.
For our database we will have kmers, and <genome name> as edges.
add_genome_schema() allows genomes to be added programmatically.
:param client: dgraph client
:return: The client altered via the schema set out here
"""
schema = """
kmer: string @index(exact, term) .
"""
try:
LOG.info("Add schema to graph with client")
return client.alter(pydgraph.Operation(schema=schema))
except:
LOG.critical("Failed to add schema to graph")
def setup_schema(self) -> None:
"""Sets up the DGraph schema based on the nodes. This inspect all attributes of all nodes,
and generates a schema for them. Each schema entry has the format `{node_type}.{field}`. If a
field is a string field, it has the `@index(exact)` predicate added to it.
An example output schema::
process.process_image string @index(exact)
process.process_id int
"""
all_node_types = inspect.getmembers(
sys.modules["beagle.nodes"],
lambda cls: inspect.isclass(cls)
and not inspect.isabstract(cls)
and issubclass(cls, Node)
and cls != Node,
)
schema = ""
for cls_name, node_class in all_node_types:
for attr, attr_type in node_class.__annotations__.items():
if attr == "key_fields":
continue
# https://github.com/python/typing/issues/528#issuecomment-357751667
if type(attr_type) == type(Union):
attr_type = attr_type.__args__[0]
if attr_type == int:
attr_type = "int"
elif type(attr_type) == type(DefaultDict) and issubclass(
attr_type.__args__[1], Edge
):
# Don't need this, get built automatically
continue
else:
attr_type = "string @index(exact)"
# Remove spaces, lowercase
schema += f"{node_class.__name__.lower().replace(' ', '_')}.{attr}: {attr_type} .\n"
schema += "<type>: string @index(exact) .\n"
logger.debug(schema)
self.dgraph.alter(pydgraph.Operation(schema=schema))
