import json
import pandas as pd
from py2neo import Graph, Node, Relationship, NodeMatcher, RelationshipMatcher 

##########################################################################
# 本代码逻辑上存在问题，观念里认为两个节点之间可以由多个标签相同而属性不同的关系，实际上是不对的；而且一般情况下每个目标服务器仅对应一个目的端口。因此本代码逻辑有问题。作废
##########################################################################


# 文件名
filename = "C:/Users/yinwe/Desktop/data.json"

# 连接数据库
graph = Graph("http://192.168.86.128:7474/", user="neo4j", password="12345678", name="neo4j")
node_matcher = NodeMatcher(graph)
relation_matcher = RelationshipMatcher(graph)


# 判断节点是否存在
def NodeExist(graph, label, ip):
    result = node_matcher.match(label).where('_.ip="'+ip+'"').first()
    if result is None:
       return False
    else:
       return True

# 判断关系是否存在
def RelationshipExist(graph, node1, node2, dst_port):
    result = relation_matcher.match({node1, node2}, r_type='Connect').where('_.dst_port="'+dst_port+'"').first()
    if result is None:
        return False
    else:
        return True


with open(filename, 'r', encoding='UTF-8') as file_in:
    for line in file_in:
        data = json.loads(line)  #json.loads用于从json文件中读取的字符串数据转为字典数据

        # 取出ip和port数据
        src_ip = data['src_ip']
        src_port = data['src_port']
        dst_ip = data['dst_ip']
        dst_port = data['dst_port']

        node1_exist = NodeExist(graph, "Src Host", src_ip)  # 源节点是否存在
        node2_exist = NodeExist(graph, "Dst Host", dst_ip)  # 目的节点是否存在

        # 判断源ip主机节点是否存在，不存在则创建，存在则不创建
        if not node1_exist:
            node1 = Node("Src Host", ip=src_ip)
            graph.create(node1)
        else:
            node1 = node_matcher.match("Src Host").where('_.ip="'+src_ip+'"').first()
            
        # 判断目的ip主机节点是否存在
        if not node2_exist:
            node2 = Node("Dst Host", ip=dst_ip)
            graph.create(node2)
        else:
            node2 = node_matcher.match("Dst Host").where('_.ip="'+dst_ip+'"').first()

        r_exist = RelationshipExist(graph, node1, node2, dst_port)  # 判断连接关系是否存在

        # 如果关系不存在，则创建关系，连接次数为1
        if not r_exist:
            r = Relationship(node1, 'Connect', node2, dst_port=dst_port, num=1)
            graph.create(r)
        # 如存在，则连接次数加1
        else:
            r = relation_matcher.match({node1, node2}, r_type='Connect').where('_.dst_port="'+dst_port+'"').first()
            r['num'] += 1
            graph.push(r)