126 lines
2.3 KiB
Go
126 lines
2.3 KiB
Go
package models
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"io/ioutil"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/segmentio/kafka-go"
|
|
"github.com/sirupsen/logrus"
|
|
)
|
|
|
|
const (
|
|
nodesFilename = "nodes.json"
|
|
edgesFilename = "edges.json"
|
|
)
|
|
|
|
// SaveData is the main function used to save data. You need to pass a specific
|
|
// flag to it and the data recieved. It just saves data, it doesn't filter or
|
|
// look for it in the stream.
|
|
func SaveData(flag string, data CertStreamStruct) {
|
|
err := fileExists(nodesFilename)
|
|
if err != nil {
|
|
logrus.Error(err)
|
|
}
|
|
|
|
err = fileExists(edgesFilename)
|
|
if err != nil {
|
|
logrus.Error(err)
|
|
}
|
|
|
|
nodeFile, err := ioutil.ReadFile(nodesFilename)
|
|
if err != nil {
|
|
logrus.Error(err)
|
|
}
|
|
|
|
edgeFile, err := ioutil.ReadFile(edgesFilename)
|
|
if err != nil {
|
|
logrus.Error(err)
|
|
}
|
|
|
|
nodeDatas := []Node{}
|
|
edgeDatas := []Edge{}
|
|
|
|
json.Unmarshal(nodeFile, &nodeDatas)
|
|
json.Unmarshal(edgeFile, &edgeDatas)
|
|
|
|
// res := CertStreamStruct{}
|
|
// json.Unmarshal(data, &res)
|
|
|
|
node := &Node{
|
|
ID: flag + "--" + uuid.New().String(),
|
|
Type: flag,
|
|
Data: data,
|
|
}
|
|
|
|
// edge := &Edge{
|
|
// ID: uuid.New(),
|
|
// NodeOneID: node.ID,
|
|
// Timestamp: time.Now(),
|
|
// }
|
|
|
|
nodeDatas = append(nodeDatas, *node)
|
|
// edgeDatas = append(edgeDatas, *edge)
|
|
|
|
nodeBytes, err := json.Marshal(nodeDatas)
|
|
if err != nil {
|
|
logrus.Error(err)
|
|
}
|
|
|
|
// edgeBytes, err := json.Marshal(edgeDatas)
|
|
// if err != nil {
|
|
// logrus.Error(err)
|
|
// }
|
|
|
|
err = ioutil.WriteFile(nodesFilename, nodeBytes, 0644)
|
|
if err != nil {
|
|
logrus.Error(err)
|
|
}
|
|
|
|
// err = ioutil.WriteFile(edgesFilename, edgeBytes, 0644)
|
|
// if err != nil {
|
|
// logrus.Error(err)
|
|
// }
|
|
}
|
|
|
|
func BuildNode(flag string, data CertStreamStruct) (*Node, error) {
|
|
err := fileExists(nodesFilename)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Node{
|
|
ID: flag + "--" + uuid.New().String(),
|
|
Type: flag,
|
|
Data: data,
|
|
}, nil
|
|
|
|
}
|
|
|
|
func SendEventToKafka(conn *kafka.Conn, node Node) error {
|
|
conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
|
_, err := conn.WriteMessages(kafka.Message{Value: []byte(fmt.Sprintf("%v", node))})
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// Helpers
|
|
|
|
func fileExists(filename string) error {
|
|
_, err := os.Stat(filename)
|
|
if os.IsNotExist(err) {
|
|
_, err := os.Create(filename)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|