def generate_user_nodes(src, dst):
"""Generate user node csv file."""
attributes = [
"uid:String",
"user_agent:String",
"device:String",
"os:String",
"browser:String",
"email:String",
"type:String",
]
with open(src) as src_data:
with gremlin_writer(GremlinNodeCSV, dst,
attributes=attributes) as writer:
for data in json_lines_file(src_data):
writer.add(
_id=data["uid"],
attribute_map={
"uid": data["uid"],
"user_agent": data["user_agent"],
"device": data["device"],
"os": data["os"],
"browser": data["browser"],
"email": data["email"],
"type": data["type"],
},
label="transientId",
)
return dst
def generate_website_csv(urls, titles, dst):
"""Generate destination CSV file."""
attributes = ["url:String", "title:String"]
with gremlin_writer(GremlinNodeCSV, dst, attributes=attributes) as writer:
for website in generate_websites(urls, titles):
attribute_map = {"url": website.url, "title": website.title}
writer.add(_id=website.url,
attribute_map=attribute_map,
label=WEBSITE_LABEL)
def generate_persistent_nodes(src, dst):
"""Generate persistent node csv file."""
with open(src) as f_h:
with gremlin_writer(GremlinNodeCSV, dst,
attributes=["pid:String"]) as writer:
for data in json_lines_file(f_h):
writer.add(
_id=data["pid"],
attribute_map={"pid": data["pid"]},
label="persistentId",
)
def generate_website_group_edges(website_group_json, dst):
"""Generate website group edges CSV."""
with open(website_group_json) as f_h:
with gremlin_writer(GremlinEdgeCSV, dst, attributes=[]) as writer:
for data in json_lines_file(f_h):
root_id = data["id"]
websites = data["websites"]
for website in websites:
writer.add(_id=get_id(root_id, website, {}),
_from=root_id,
to=website,
label=WEBISTE_GROUP_EDGE_LABEL,
attribute_map={})
def generate_persistent_id_edges(src, dst):
"""Generate persistentID edges based on union-find datastructure."""
with open(src) as f_h:
with gremlin_writer(GremlinEdgeCSV, dst, attributes=[]) as writer:
for data in json_lines_file(f_h):
for node in data["transientIds"]:
persistent_to_transient = {
"_id": get_id(data["pid"], node, {}),
"_from": data["pid"],
"to": node,
"label": "has_identity",
"attribute_map": {},
}
writer.add(**persistent_to_transient)
def generate_website_group_nodes(website_group_json, dst):
"""Generate website groups csv."""
attributes = ["url:String", "category:String", "categoryCode:String"]
with open(website_group_json) as f_h:
with gremlin_writer(GremlinNodeCSV, dst,
attributes=attributes) as writer:
for data in json_lines_file(f_h):
writer.add(_id=data["id"],
attribute_map={
"url": data["url"],
"category": data["category"]["name"],
"categoryCode": data["category"]["code"]
},
label=WEBSITE_GROUP_LABEL)
def generate_identity_group_nodes(src, dst):
"""Generate identity_group csv file with nodes."""
attrs = ["igid:String", "type:String"]
with open(src) as f_h:
with gremlin_writer(GremlinNodeCSV, dst, attributes=attrs) as writer:
for data in json_lines_file(f_h):
if data["persistentIds"]:
writer.add(
_id=data["igid"],
attribute_map={
"igid": data["igid"],
"type": data["type"]
},
label="identityGroup",
)
def generate_identity_group_edges(src, dst):
"""Generate identity_group edge csv file."""
with open(src) as f_h:
with gremlin_writer(GremlinEdgeCSV, dst, attributes=[]) as writer:
for data in json_lines_file(f_h):
persistent_ids = data["persistentIds"]
if persistent_ids:
for persistent_id in persistent_ids:
identity_group_to_persistent = {
"_id": get_id(data["igid"], persistent_id, {}),
"_from": data["igid"],
"to": persistent_id,
"attribute_map": {},
"label": "member",
}
writer.add(**identity_group_to_persistent)
def generate_user_website_edges(src_map, dst):
"""Generate edges between user nodes and website nodes."""
with open(src_map["urls"]) as url_file:
fact_to_website = {}
for row in csv.reader(url_file, delimiter=","):
fact_to_website[int(row[0])] = row[1]
with open(src_map["facts"]) as facts_file:
attrs = [
"ts:Date",
"visited_url:String",
"uid:String",
"state:String",
"city:String",
"ip_address:String",
]
with gremlin_writer(GremlinEdgeCSV, dst, attributes=attrs) as writer:
for data in json_lines_file(facts_file):
for fact in data["facts"]:
timestamp = _parse_ts(fact["ts"])
website_id = fact_to_website[fact["fid"]]
loc_attrs = {
"state": fact["state"],
"city": fact["city"],
"ip_address": fact["ip_address"],
}
attr_map = {
"ts": timestamp,
"visited_url": website_id,
"uid": data["uid"],
**loc_attrs,
}
user_to_website = {
"_id": get_id(data["uid"], website_id, attr_map),
"_from": data["uid"],
"to": website_id,
"label": "visited",
"attribute_map": attr_map,
}
try:
writer.add(**user_to_website)
except Exception:
logger.exception("Something went wrong while creating an edge")
logger.info(json.dumps({"uid": data["uid"], **fact}))
return dst
def generate_ip_loc_edges_from_facts(src, dst):
"""Generate ip location csv file with edges."""
with open(src) as f_h:
with gremlin_writer(GremlinEdgeCSV, dst, attributes=[]) as writer:
for data in json_lines_file(f_h):
uid_locations = set()
for fact in data["facts"]:
uid_locations.add(
IPLoc(fact["state"], fact["city"], fact["ip_address"]))
for location in uid_locations:
loc_id = get_id(location)
writer.add(
_id=get_edge_id(data["uid"], loc_id, {}),
_from=data["uid"],
to=loc_id,
label="uses",
attribute_map={},
)
def generate_ip_loc_nodes_from_facts(src, dst):
"""Generate ip location csv file with nodes."""
attrs = ["state:String", "city:String", "ip_address:String"]
with open(src) as f_h:
with gremlin_writer(GremlinNodeCSV, dst, attributes=attrs) as writer:
locations = set()
for data in json_lines_file(f_h):
for fact in data["facts"]:
locations.add(
IPLoc(fact["state"], fact["city"], fact["ip_address"]))
for location in locations:
writer.add(
_id=get_id(location),
attribute_map={
"state": location.state,
"city": location.city,
"ip_address": location.ip_address,
},
label="IP",
)