-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmain.py
169 lines (106 loc) · 4.95 KB
/
main.py
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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import re
import zipfile
import os
import sys
import csv
import itertools
from os import listdir
from os.path import isfile, join
from xml.dom.minidom import parse, parseString
from pyspark import SparkContext
from pyspark import SparkConf
from pyspark.sql import SQLContext
from neo4j.v1 import GraphDatabase, basic_auth
sc = SparkContext("local", "Enron")
sqlContext = SQLContext(sc)
directory_to_source = 'enron_db'
directory_to_extract_to = '/xml'
def extractXML(path_to_zip_file):
full_path_to_zip_file = join(directory_to_source,path_to_zip_file)
zip_ref = zipfile.ZipFile(full_path_to_zip_file, 'r')
log = open("xml/xml_names.txt" , "a")
for f in zip_ref.namelist():
if f.startswith('zl'):
zip_ref.extract(f, directory_to_extract_to)
log.writelines(f + "\n")
log.close()
def getEmails(document, item):
emails = []
to = []
frm = []
tags = document.getElementsByTagName('Tag')
for t in tags:
if t.getAttribute('TagName') == '#To':
l = t.getAttribute('TagValue')
to = re.findall(r'[\w\.-]+@[\w\.-]+', l)
if t.getAttribute('TagName') == '#From':
l = t.getAttribute('TagValue')
frm = re.findall(r'[\w\.-]+@[\w\.-]+', l)
if to:
if frm:
emails.extend(itertools.product(to, frm))
return emails
def xml_to_emails(item):
emailed = []
doc = parse(join(directory_to_extract_to,item))
documents = doc.getElementsByTagName('Document') #pega todos os documentos (emails) indexados no xml
for document in documents:
res1 = getEmails(document, item) #extrai destinatarios de cada email
if res1:
emailed.extend(res1)
return emailed
def generateEmailRDD():
with open("xml/xml_names.txt") as f:
content = f.readlines()
content = [x.strip() for x in content]
raw_xml = sc.parallelize(content)
rdd = raw_xml.flatMap(lambda x: xml_to_emails(x)).map(lambda word: (word, 1)).reduceByKey(lambda a, b: a + b)
return rdd
def edge_format_csv(f):
return [(f[0][0],f[0][1],f[1])]
def node_format_csv(f):
return [(f[0][0],f[0][0]),(f[0][1],f[0][1])]
def generate_neo4j_graph(node_csv_name, edge_csv_name):
print("Building graph...")
driver = GraphDatabase.driver("bolt://127.0.0.1", auth=("neo4j", "showdotomas"),encrypted=False)
session = driver.session()
clean = """MATCH (n)
DETACH DELETE n
"""
session.run(clean)
load = """
LOAD CSV WITH HEADERS FROM "file:///nodes/"""+node_csv_name+"""" AS csvLine
CREATE (e:Employee {username: csvLine.name})
"""
session.run(load)
edge = """
USING PERIODIC COMMIT 500
LOAD CSV WITH HEADERS FROM "file:///edges/"""+edge_csv_name+"""" AS csvLine
MATCH (src:Employee),(dst:Employee)
WHERE src.username =csvLine.src AND dst.username=csvLine.dst
CREATE (src)-[:EMAIL { weight: toInt(csvLine.weight) }]->(dst)
"""
session.run(edge)
print("Completed.")
session.close()
if __name__ == '__main__':
if len(sys.argv) == 2:
zfiles = sys.argv[1]
else:
zfiles = [f for f in os.listdir(directory_to_source) if f.endswith('xml.zip')] #lista com todos os zip (um pra cada remetente de email)
open("xml/xml_names.txt", 'w').close() #zera arquivo
for zfile in zfiles: #extrair todos os xml para pasta
extractXML(zfile)
rdd = generateEmailRDD()
edge_rdd = rdd.flatMap(lambda x: edge_format_csv(x)).toDF(["src","dst","weight"])
node_rdd = rdd.flatMap(lambda x: node_format_csv(x)).distinct().toDF(["id","name"])
print("Numer of nodes: " +str(node_rdd.count()))
print("Numer of edges: " +str(edge_rdd.count()))
has_csv = [f for f in os.listdir(r'/home/lucas/.config/Neo4j Desktop/Application/neo4jDatabases/database-9be871ed-d8ad-4ec3-b24a-71b5431b0fbd/current/import') if f.endswith('es')]
if len(has_csv)==0:
edge_rdd.coalesce(1).write.format('com.databricks.spark.csv').options(header='true').save(r'/home/lucas/.config/Neo4j Desktop/Application/neo4jDatabases/database-9be871ed-d8ad-4ec3-b24a-71b5431b0fbd/current/import/edges')
node_rdd.coalesce(1).write.format('com.databricks.spark.csv').options(header='true').save(r'/home/lucas/.config/Neo4j Desktop/Application/neo4jDatabases/database-9be871ed-d8ad-4ec3-b24a-71b5431b0fbd/current/import/nodes')
node_csv_name = [f for f in os.listdir(r'/home/lucas/.config/Neo4j Desktop/Application/neo4jDatabases/database-9be871ed-d8ad-4ec3-b24a-71b5431b0fbd/current/import/nodes') if f.endswith('.csv')]
edge_csv_name = [f for f in os.listdir(r'/home/lucas/.config/Neo4j Desktop/Application/neo4jDatabases/database-9be871ed-d8ad-4ec3-b24a-71b5431b0fbd/current/import/edges') if f.endswith('.csv')]
if node_csv_name and edge_csv_name:
generate_neo4j_graph(node_csv_name[0], edge_csv_name[0])