styx/models/main.go
Christopher Talib f61fe566a5 Basic connection to Dgraph DB
The first work and input to the graph db is set up in this work. It's
for the moment very basic and doesn't cover relations and only works for
certstream data.
2020-03-04 15:16:59 +01:00

425 lines
11 KiB
Go

package models
import (
"encoding/json"
"io/ioutil"
"time"
"github.com/google/uuid"
"github.com/ns3777k/go-shodan/v4/shodan"
"github.com/sirupsen/logrus"
"gitlab.dcso.lolcat/LABS/styx/balboa"
"gitlab.dcso.lolcat/LABS/styx/utils"
)
/**
Structure of this file:
* type
* build node functions
* save node functions
**/
// Node defines the data we gather through the parsing. It should follow the
// Styx terminology
// (https://docs.google.com/document/d/1dIrh1Lp3KAjEMm8o2VzAmuV0Peu-jt9aAh1IHrjAroM/pub#h.xzbicbtscatx)
type Node struct {
UID string `json:"uid,omiempty"`
ID string `json:"id,omiempty"`
Type string `json:"type,omiempty"`
Data string `json:"data,omiempty"`
Created string `json:"created,omiempty"`
Modified string `json:"modified,omiempty"`
}
// BuildNode builds a node to send to MQ instance.
func BuildNode(flag string, dataType string, data string) *Node {
t := time.Now()
rfc3339time := t.Format(time.RFC3339)
uuid := uuid.New().String()
return &Node{
UID: "_:" + flag + "--" + uuid,
ID: flag + "--" + uuid,
Type: dataType,
Data: data,
Created: rfc3339time,
Modified: rfc3339time,
}
}
// SaveNode saves a node to a file.
func SaveNode(filename string, node *Node) {
err := utils.FileExists(filename)
if err != nil {
logrus.Error(err)
}
nodeFile, err := ioutil.ReadFile(filename)
if err != nil {
logrus.Error(err)
}
nodeDatas := []Node{}
if err := json.Unmarshal(nodeFile, &nodeDatas); err != nil {
logrus.Error(err)
}
nodeDatas = append(nodeDatas, *node)
nodeBytes, err := json.Marshal(nodeDatas)
if err != nil {
logrus.Error(err)
}
err = ioutil.WriteFile(filename, nodeBytes, 0644)
if err != nil {
logrus.Error(err)
}
}
// Edge defines a relation between two nodes.
type Edge struct {
UID string `json:"uid,omiempty"`
ID string `json:"id,omiempty"`
NodeOneID string `json:"nodeOneID,omiempty"`
NodeTwoID string `json:"nodeTwoID,omiempty"`
Timestamp string `json:"timestamp,omiempty"`
Source string `json:"source,omiempty"`
}
// BuildEdge build a send from two nodes with a given source type.
func BuildEdge(source string, nodeOneUUID string, nodeTwoUUID string) *Edge {
t := time.Now()
rfc3339time := t.Format(time.RFC3339)
uuid := uuid.New().String()
return &Edge{
UID: "_:" + "edge--" + uuid,
ID: "edge--" + uuid,
Source: source,
NodeOneID: nodeOneUUID,
NodeTwoID: nodeTwoUUID,
Timestamp: rfc3339time,
}
}
// SaveEdge saves an edge to a file.
func SaveEdge(edge *Edge) {
err := utils.FileExists("edges.json")
if err != nil {
logrus.Error(err)
}
edgeFile, err := ioutil.ReadFile("edges.json")
if err != nil {
logrus.Error(err)
}
edgeDatas := []Edge{}
if err := json.Unmarshal(edgeFile, &edgeDatas); err != nil {
logrus.Error(err)
}
edgeDatas = append(edgeDatas, *edge)
edgeBytes, err := json.Marshal(edgeDatas)
if err != nil {
logrus.Error(err)
}
err = ioutil.WriteFile("edges.json", edgeBytes, 0644)
if err != nil {
logrus.Error(err)
}
}
// CertStreamRaw is a wrapper around the stream function to unmarshall the
// data receive in a Go structure.
type CertStreamRaw struct {
UID string `json:"uid,omiempty"`
ID string `json:"id,omiempty"`
Type string `json:"type,omiempty"`
Data CertStreamStruct `json:"data,omiempty"`
Created string `json:"created,omiempty"`
Modified string `json:"modified,omiempty"`
}
// CertNode represents our custom struct of data extraction from CertStream.
type CertNode struct {
UID string `json:"uid,omiempty"`
ID string `json:"id,omiempty"`
Fingerprint string `json:"fingerprint,omiempty"`
NotBefore string `json:"notBefore,omiempty"`
NotAfter string `json:"notAfter,omiempty"`
CN string `json:"cn,omiempty"`
SourceName string `json:"sourceName,omiempty"`
SerialNumber string `json:"serialNumber,omiempty"`
BasicConstraints string `json:"basicConstraints,omiempty"`
RawUUID string `json:"rawUUID,omiempty"`
Chain []CertNode `json:"chainedTo,omiempty"`
}
// WrapCertStreamData is a wrapper around CertStreamStruct.
func WrapCertStreamData(data CertStreamStruct) *CertStreamRaw {
t := time.Now()
rfc3339time := t.Format(time.RFC3339)
return &CertStreamRaw{
UID: "_:certstream--" + uuid.New().String(),
ID: "certstream--" + uuid.New().String(),
Type: "certstream_raw",
Data: data,
Created: rfc3339time,
Modified: rfc3339time,
}
}
// BuildCertNode builds a custom node based on CertStream.
func BuildCertNode(rawNode *CertStreamRaw) *CertNode {
main := &CertNode{
UID: "_:certstream--" + uuid.New().String(),
ID: "certstream--" + uuid.New().String(),
Fingerprint: rawNode.Data.Data.LeafCert.Fingerprint,
NotBefore: time.Unix(int64(rawNode.Data.Data.LeafCert.NotBefore), 0).Format(time.RFC3339),
NotAfter: time.Unix(int64(rawNode.Data.Data.LeafCert.NotAfter), 0).Format(time.RFC3339),
CN: rawNode.Data.Data.LeafCert.Subject.CN,
SourceName: rawNode.Data.Data.Source.Name,
BasicConstraints: rawNode.Data.Data.LeafCert.Extensions.BasicConstrains,
RawUUID: rawNode.ID,
}
var res []CertNode
if len(rawNode.Data.Data.Chain) > 0 {
chain := CertNode{
UID: "_:certstream--" + uuid.New().String(),
ID: "certstream--" + uuid.New().String(),
Fingerprint: rawNode.Data.Data.LeafCert.Fingerprint,
NotBefore: time.Unix(int64(rawNode.Data.Data.LeafCert.NotBefore), 0).Format(time.RFC3339),
NotAfter: time.Unix(int64(rawNode.Data.Data.LeafCert.NotAfter), 0).Format(time.RFC3339),
CN: rawNode.Data.Data.LeafCert.Subject.CN,
SourceName: rawNode.Data.Data.Source.Name,
BasicConstraints: rawNode.Data.Data.LeafCert.Extensions.BasicConstrains,
RawUUID: rawNode.ID,
}
res = append(res, chain)
}
main.Chain = res
return main
}
// SaveCertStreamRaw save the raw CertStream data.
func SaveCertStreamRaw(filename string, data *CertStreamRaw) {
err := utils.FileExists(filename)
if err != nil {
logrus.Error(err)
}
nodeFile, err := ioutil.ReadFile(filename)
if err != nil {
logrus.Error(err)
}
rawDatas := []CertStreamRaw{}
if err := json.Unmarshal(nodeFile, &rawDatas); err != nil {
logrus.Error(err)
}
rawDatas = append(rawDatas, *data)
rawBytes, err := json.Marshal(rawDatas)
if err != nil {
logrus.Error(err)
}
err = ioutil.WriteFile(filename, rawBytes, 0644)
if err != nil {
logrus.Error(err)
}
}
// SaveCertNode saves a CertNode to a json file.
func SaveCertNode(filename string, node *CertNode) {
err := utils.FileExists(filename)
if err != nil {
logrus.Error(err)
}
nodeFile, err := ioutil.ReadFile(filename)
if err != nil {
logrus.Error(err)
}
nodeDatas := []CertNode{}
if err := json.Unmarshal(nodeFile, &nodeDatas); err != nil {
logrus.Error(err)
}
nodeDatas = append(nodeDatas, *node)
nodeBytes, err := json.Marshal(nodeDatas)
if err != nil {
logrus.Error(err)
}
err = ioutil.WriteFile(filename, nodeBytes, 0644)
if err != nil {
logrus.Error(err)
}
}
// PasteNode is a node from PasteBin.
type PasteNode struct {
ID string `json:"id"`
Type string `json:"type"`
Data FullPaste `json:"data"`
Created string `json:"create"`
Modified string `json:"modified"`
}
// FullPaste wrapes meta and information from Pastebin.
type FullPaste struct {
Meta PasteMeta `json:"meta"`
Full string `json:"full"`
}
// BuildPasteNode builds a node from a FullPaste data.
func BuildPasteNode(data *FullPaste) *PasteNode {
t := time.Now()
rfc3339time := t.Format(time.RFC3339)
return &PasteNode{
ID: "pastebin--" + uuid.New().String(),
Type: "pastebin",
Data: *data,
Created: rfc3339time,
Modified: rfc3339time,
}
}
// SavePaste saves a object received from PasteBin.
func SavePaste(filename string, data *PasteNode) {
err := utils.FileExists(filename)
if err != nil {
logrus.Error(err)
}
pasteFile, err := ioutil.ReadFile(filename)
if err != nil {
logrus.Error(err)
}
rawPaste := []PasteNode{}
if err := json.Unmarshal(pasteFile, &rawPaste); err != nil {
logrus.Error(err)
}
rawPaste = append(rawPaste, *data)
rawBytes, err := json.Marshal(rawPaste)
if err != nil {
logrus.Error(err)
}
err = ioutil.WriteFile(filename, rawBytes, 0644)
if err != nil {
logrus.Error(err)
}
}
// ShodanNode is node around the shodan.HostData struct.
type ShodanNode struct {
ID string `json:"id"`
Type string `json:"type"`
Data *shodan.HostData `json:"data"`
Created string `json:"created"`
Modified string `json:"modified"`
}
// BuildShodanNode builds a wrapper node around shodan.HostData.
func BuildShodanNode(data *shodan.HostData) *ShodanNode {
t := time.Now()
rfc3339time := t.Format(time.RFC3339)
return &ShodanNode{
ID: "shodan--" + uuid.New().String(),
Type: "shodan_stream",
Data: data,
Created: rfc3339time,
Modified: rfc3339time,
}
}
// SaveShodanNode saves the raw nodes from Shodan.
func SaveShodanNode(filename string, data *ShodanNode) {
err := utils.FileExists(filename)
if err != nil {
logrus.Error(err)
}
nodeFile, err := ioutil.ReadFile(filename)
if err != nil {
logrus.Error(err)
}
rawDatas := []ShodanNode{}
if err := json.Unmarshal(nodeFile, &rawDatas); err != nil {
logrus.Error(err)
}
rawDatas = append(rawDatas, *data)
rawBytes, err := json.Marshal(rawDatas)
if err != nil {
logrus.Error(err)
}
err = ioutil.WriteFile(filename, rawBytes, 0644)
if err != nil {
logrus.Error(err)
}
}
// BalboaNode represents a return from Balboa.
type BalboaNode struct {
ID string `json:"id"`
Type string `json:"type"`
Data []balboa.Entries `json:"data"`
Created string `json:"created"`
Modified string `json:"modified"`
}
// BuildBalboaNode builds a node coming from Balboa resolution.
func BuildBalboaNode(data []balboa.Entries) *BalboaNode {
t := time.Now()
rfc3339time := t.Format(time.RFC3339)
return &BalboaNode{
ID: "balboa--" + uuid.New().String(),
Type: "balboa",
Data: data,
Created: rfc3339time,
Modified: rfc3339time,
}
}
// SaveBalboaNode saves a Balboa node.
func SaveBalboaNode(filename string, data *BalboaNode) {
err := utils.FileExists(filename)
if err != nil {
logrus.Error(err)
}
nodeFile, err := ioutil.ReadFile(filename)
if err != nil {
logrus.Error(err)
}
rawDatas := []BalboaNode{}
if err := json.Unmarshal(nodeFile, &rawDatas); err != nil {
logrus.Error(err)
}
rawDatas = append(rawDatas, *data)
rawBytes, err := json.Marshal(rawDatas)
if err != nil {
logrus.Error(err)
}
err = ioutil.WriteFile(filename, rawBytes, 0644)
if err != nil {
logrus.Error(err)
}
}