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.

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/DataDog/zstd v1.4.4 // indirect
 	github.com/cncf/udpa/go v0.0.0-20200124205748-db4b343e48c1 // indirect
 	github.com/dgraph-io/dgo v1.0.0
+	github.com/dgraph-io/dgo/v2 v2.2.0
 	github.com/envoyproxy/go-control-plane v0.9.4 // indirect
 	github.com/frankban/quicktest v1.7.2 // indirect
 	github.com/golang/mock v1.4.1 // indirect

go.sum

@@ -56,6 +56,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgraph-io/dgo v1.0.0 h1:DRuI66G+j0XWDOXly4v5PSk2dGkbIopAZIirRjq7lzI=
 github.com/dgraph-io/dgo v1.0.0/go.mod h1:6K5zUB6Lsml4SEStX+fPzGhJtCLX9XxbkHJLsGOXS1E=
+github.com/dgraph-io/dgo/v2 v2.2.0 h1:qYbm6mEF3wuKiRpgNOldk6PmPbBJFwj6vL7I7dTSdyc=
+github.com/dgraph-io/dgo/v2 v2.2.0/go.mod h1:LJCkLxm5fUMcU+yb8gHFjHt7ChgNuz3YnQQ6MQkmscI=
 github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumCAMpl/TFQ4/5kLM=
 github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
 github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no=
@@ -84,6 +86,7 @@ github.com/go-sql-driver/mysql v1.4.1/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
+github.com/gogo/protobuf v1.3.1/go.mod h1:SlYgWuQ5SjCEi6WLHjHCa1yvBfUnHcTbrrZtXPKa29o=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
 github.com/golang/groupcache v0.0.0-20190129154638-5b532d6fd5ef/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
@@ -146,6 +149,7 @@ github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7
 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
 github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
 github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvWXihfKN4Q=
+github.com/kisielk/errcheck v1.2.0/go.mod h1:/BMXB+zMLi60iA8Vv6Ksmxu/1UDYcXs4uQLJ+jE2L00=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1 h1:mweAR1A6xJ3oS2pRaGiHgQ4OO8tzTaLawm8vnODuwDk=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
@@ -337,6 +341,7 @@ golang.org/x/sys v0.0.0-20190507160741-ecd444e8653b/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20190606165138-5da285871e9c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190624142023-c5567b49c5d0/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190726091711-fc99dfbffb4e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20190801041406-cbf593c0f2f3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191001151750-bb3f8db39f24/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191204072324-ce4227a45e2e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191228213918-04cbcbbfeed8/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
@@ -358,6 +363,7 @@ golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxb
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/tools v0.0.0-20181030221726-6c7e314b6563/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=

graph/main.go

@@ -3,118 +3,119 @@ package graph
 import (
 	"context"
 
-	"github.com/dgraph-io/dgo"
+	"github.com/dgraph-io/dgo/v2"
-	"github.com/dgraph-io/dgo/protos/api"
+	"github.com/dgraph-io/dgo/v2/protos/api"
 	"google.golang.org/grpc"
 )
 
-func ConnectToDgraph() error {
+func ConnectToDgraph() (*dgo.Dgraph, error) {
 	conn, err := grpc.Dial("localhost:9080", grpc.WithInsecure())
 	if err != nil {
-		return err
+		return nil, err
 	}
-	defer conn.Close()
 
 	dgraphClient := dgo.NewDgraphClient(api.NewDgraphClient(conn))
 
 	err = setupDgraphSchema(dgraphClient)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	return nil
+	return dgraphClient, nil
 }
 
 func setupDgraphSchema(c *dgo.Dgraph) error {
 	err := c.Alter(context.Background(), &api.Operation{
 		DropAll: true,
+	})
+
+	err = c.Alter(context.Background(), &api.Operation{
 		Schema: `
-ID: string @index(term) .
+id: string @index(term) .
-Type: string @index(term) .
+type: string @index(term) .
-Data: string .
+data: string .
-NodeOneID: [uid] .
+nodeOneID: uid .
-NodeTwoID: [uid] .
+nodeTwoID: uid .
-Source: string @index(term) .
+source: string @index(term) .
-Timestamp: string .
+timestamp: string .
-Created: dateTime .
+created: dateTime .
-Modified: dateTime .
+modified: dateTime .
 
 type Node {
-ID
+id
-Type
+type
-Data
+data
-Created
+created
-Modified
+modified
 }
 
 type Edge {
-ID
+id
-NodeOneID
+nodeOneID
-NodeTwoID
+nodeTwoID
-Timestamp
+timestamp
-Source
+source
 }
 
-Fingerprint: string .
+fingerprint: string .
-NotBefore: string .
+notBefore: string .
-NotAfter: string .
+notAfter: string .
-CN: string .
+cn: string .
-SourceName: string .
+sourceName: string .
-SerialNumber: string .
+serialNumber: string .
-BasicConstraints: string .
+basicConstraints: string .
-Chain: [uid] .
+chain: uid .
-Data: [uid] .
 
 type CertNode {
-ID
+id
-Fingerprint
+fingerprint
-NotBefore
+notBefore
-NotAfter
+notAfter
-CN
+cn
-SourceName
+sourceName
-SerialNumber
+serialNumber
-BasicConstraints
+basicConstraints
-Chain
+chain
 }
 
 type CertRaw {
-ID
+id
-Type
+type
-Created
+created
-Modified
+modified
-Data
+data
 }
 
 type PasteNode {
-ID
+id
-Type
+type
-Created
+created
-Modified
+modified
-Data
+data
 }
 
-Meta: [uid] .
+meta: uid .
-Full: string .
+full: string .
 
 type FullPaste {
-Meta
+meta
-Full
+full
 }
 
 type ShodanNode {
-ID
+id
-Type
+type
-Data
+data
-Created
+created
-Modified
+modified
 }
 
 type BalboaNode {
-ID
+id
-Type
+type
-Data
+data
-Created
+created
-Modified
+modified
 }
 `})
 	if err != nil {

main.go

@@ -36,7 +36,7 @@ func main() {
 	// panic(err)
 	// }
 
-	err = graph.ConnectToDgraph()
+	dgraphClient, err := graph.ConnectToDgraph()
 	if err != nil {
 		logrus.WithField("err", err).Error("error initialising the graph database")
 	}
@@ -52,7 +52,7 @@ func main() {
 	if ok := c.Initialize(); !ok {
 		logrus.Info("certstream plugin not activated")
 	}
-	c.Run(&wg)
+	c.Run(&wg, dgraphClient)
 
 	// pastebin
 	p := plugins.PastebinPlugin{}

models/main.go

@@ -25,7 +25,7 @@ Structure of this file:
 // Styx terminology
 // (https://docs.google.com/document/d/1dIrh1Lp3KAjEMm8o2VzAmuV0Peu-jt9aAh1IHrjAroM/pub#h.xzbicbtscatx)
 type Node struct {
-	Uid      string `json:"uid,omiempty"`
+	UID      string `json:"uid,omiempty"`
 	ID       string `json:"id,omiempty"`
 	Type     string `json:"type,omiempty"`
 	Data     string `json:"data,omiempty"`
@@ -37,9 +37,10 @@ type Node struct {
 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.New().String(),
+		UID:      "_:" + flag + "--" + uuid,
-		ID:       flag + "--" + uuid.New().String(),
+		ID:       flag + "--" + uuid,
 		Type:     dataType,
 		Data:     data,
 		Created:  rfc3339time,
@@ -78,7 +79,7 @@ func SaveNode(filename string, node *Node) {
 
 // Edge defines a relation between two nodes.
 type Edge struct {
-	Uid       string `json:"uid,omiempty"`
+	UID       string `json:"uid,omiempty"`
 	ID        string `json:"id,omiempty"`
 	NodeOneID string `json:"nodeOneID,omiempty"`
 	NodeTwoID string `json:"nodeTwoID,omiempty"`
@@ -90,9 +91,10 @@ type Edge struct {
 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.New().String(),
+		UID:       "_:" + "edge--" + uuid,
-		ID:        "edge--" + uuid.New().String(),
+		ID:        "edge--" + uuid,
 		Source:    source,
 		NodeOneID: nodeOneUUID,
 		NodeTwoID: nodeTwoUUID,
@@ -132,16 +134,17 @@ func SaveEdge(edge *Edge) {
 // CertStreamRaw is a wrapper around the stream function to unmarshall the
 // data receive in a Go structure.
 type CertStreamRaw struct {
-	ID       string           `json:"id"`
+	UID      string           `json:"uid,omiempty"`
-	Type     string           `json:"type"`
+	ID       string           `json:"id,omiempty"`
-	Data     CertStreamStruct `json:"data"`
+	Type     string           `json:"type,omiempty"`
-	Created  string           `json:"created"`
+	Data     CertStreamStruct `json:"data,omiempty"`
-	Modified string           `json:"modified"`
+	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"`
+	UID              string     `json:"uid,omiempty"`
 	ID               string     `json:"id,omiempty"`
 	Fingerprint      string     `json:"fingerprint,omiempty"`
 	NotBefore        string     `json:"notBefore,omiempty"`
@@ -159,6 +162,7 @@ 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,
@@ -170,6 +174,7 @@ func WrapCertStreamData(data CertStreamStruct) *CertStreamRaw {
 // 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),
@@ -183,6 +188,7 @@ func BuildCertNode(rawNode *CertStreamRaw) *CertNode {
 	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),

plugins/certstream.go

@@ -1,9 +1,14 @@
 package plugins
 
 import (
+	"context"
+	"encoding/json"
+	"fmt"
 	"sync"
 
 	"github.com/CaliDog/certstream-go"
+	"github.com/dgraph-io/dgo/v2"
+	"github.com/dgraph-io/dgo/v2/protos/api"
 	"github.com/jmoiron/jsonq"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"
@@ -34,11 +39,11 @@ func (c *CertStreamPlugin) Initialize() bool {
 }
 
 // Run runs the Certstream plugin.
-func (c *CertStreamPlugin) Run(wg *sync.WaitGroup) {
+func (c *CertStreamPlugin) Run(wg *sync.WaitGroup, dgraphClient *dgo.Dgraph) {
 	if !c.Running {
 		c.StopChan = make(chan bool)
 		wg.Add(1)
-		go c.doRun()
+		go c.doRun(dgraphClient)
 		c.Running = true
 	}
 }
@@ -54,7 +59,7 @@ func (c *CertStreamPlugin) Stop(wg *sync.WaitGroup) {
 	}
 }
 
-func (c *CertStreamPlugin) doRun() {
+func (c *CertStreamPlugin) doRun(graphClient *dgo.Dgraph) {
 	for {
 		select {
 		case jq := <-c.Stream:
@@ -75,6 +80,36 @@ func (c *CertStreamPlugin) doRun() {
 						edge = models.BuildEdge("certstream", mainNode.ID, certNode.ID)
 						models.SaveEdge(edge)
 						// saveSingleValues(conn, "certstream", "domain", certNode.ID, domain)
+
+						mu := &api.Mutation{
+							CommitNow: true,
+						}
+						marshaled, err := json.Marshal(mainNode)
+						if err != nil {
+							logrus.Fatal(err)
+						}
+						mu.SetJson = marshaled
+						assigned, err := graphClient.NewTxn().Mutate(context.Background(), mu)
+						if err != nil {
+							logrus.Fatal(err)
+						}
+						variables := map[string]string{"$id": assigned.Uids[mainNode.ID]}
+						q := `query Node($id: string){
+node(func: uid($id)) {
+uid
+ID
+Type
+Data
+Created
+Modified
+}
+}`
+						resp, err := graphClient.NewTxn().QueryWithVars(context.Background(), q, variables)
+						if err != nil {
+							logrus.Fatal(err)
+						}
+						fmt.Println(resp)
+
 					}
 				}