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,omitempty"`
	ID         string     `json:"id,omitempty"`
	NodeType   string     `json:"nodeType,omitempty"`
	NData      string     `json:"ndata,omitempty"`
	Created    string     `json:"created,omitempty"`
	Modified   string     `json:"modified,omitempty"`
	DType      []string   `json:"dgraph.type,omitempty"`
	CertNode   CertNode   `json:"certNode,omitempty"`
	ShodanNode ShodanNode `json:"shodanNode,omitempty"`
	PasteNode  PasteNode  `json:"pasteNode,omitempty"`
	Match      Match      `json:"match,omitempty"`
}

// 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{
		ID:       flag + "--" + uuid,
		NodeType: dataType,
		NData:    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 {
	ID        string                 `json:"id,omitempty"`
	NodeOne   map[string]interface{} `json:"nodeOne,omitempty"`
	NodeTwo   map[string]interface{} `json:"nodeTwo,omitempty"`
	Timestamp string                 `json:"timestamp,omitempty"`
	Source    string                 `json:"source,omitempty"`
	DType     []string               `json:"dgraph.type,omitempty"`
}

// BuildEdge build a send from two nodes with a given source type.
func BuildEdge(source string, nodeOne, nodeTwo map[string]interface{}) *Edge {
	t := time.Now()
	rfc3339time := t.Format(time.RFC3339)
	uuid := uuid.New().String()
	return &Edge{
		ID:        "edge--" + uuid,
		Source:    source,
		NodeOne:   nodeOne,
		NodeTwo:   nodeTwo,
		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)
	}
}

// Match represents clustered results based on a target.
type Match struct {
	ID        string `json:"id,omitempty"`
	UID       string `json:"uid,omitempty"`
	Nodes     []Node `json:"nodes,omitempty"`
	Target    string `json:"target,omitempty"`
	Timestamp string `json:"timestamp,omitempty"`
	NodeType  string `json:"nodeType,omitempty"`
}

// CertStreamRaw is a wrapper around the stream function to unmarshall the
// data receive in a Go structure.
type CertStreamRaw struct {
	ID       string           `json:"id,omitempty"`
	NodeType string           `json:"nodeType,omitempty"`
	Data     CertStreamStruct `json:"data,omitempty"`
	Created  string           `json:"created,omitempty"`
	Modified string           `json:"modified,omitempty"`
}

// CertNode represents our custom struct of data extraction from CertStream.
type CertNode struct {
	ID               string        `json:"id,omitempty"`
	Fingerprint      string        `json:"fingerprint,omitempty"`
	NotBefore        string        `json:"notBefore,omitempty"`
	NotAfter         string        `json:"notAfter,omitempty"`
	CN               string        `json:"cn,omitempty"`
	SourceName       string        `json:"sourceName,omitempty"`
	SerialNumber     string        `json:"serialNumber,omitempty"`
	BasicConstraints string        `json:"basicConstraints,omitempty"`
	Raw              CertStreamRaw `json:"raw,omitempty"`
	Chain            []CertNode    `json:"chain,omitempty"`
}

// WrapCertStreamData is a wrapper around CertStreamStruct.
func WrapCertStreamData(data CertStreamStruct) *CertStreamRaw {
	t := time.Now()
	rfc3339time := t.Format(time.RFC3339)
	return &CertStreamRaw{
		ID:       "certstream--" + uuid.New().String(),
		NodeType: "certstream_raw",
		Data:     data,
		Created:  rfc3339time,
		Modified: rfc3339time,
	}
}

// BuildCertNode builds a custom node based on CertStream.
func BuildCertNode(rawNode *CertStreamRaw) *CertNode {
	main := &CertNode{
		ID:               "certstream--" + uuid.New().String(),
		Fingerprint:      rawNode.Data.CSData.LeafCert.Fingerprint,
		NotBefore:        time.Unix(int64(rawNode.Data.CSData.LeafCert.NotBefore), 0).Format(time.RFC3339),
		NotAfter:         time.Unix(int64(rawNode.Data.CSData.LeafCert.NotAfter), 0).Format(time.RFC3339),
		CN:               rawNode.Data.CSData.LeafCert.Subject.CN,
		SourceName:       rawNode.Data.CSData.Source.Name,
		BasicConstraints: rawNode.Data.CSData.LeafCert.Extensions.BasicConstrains,
		Raw:              *rawNode,
	}

	var res []CertNode
	if len(rawNode.Data.CSData.Chain) > 0 {
		chain := CertNode{
			ID:               "certstream--" + uuid.New().String(),
			Fingerprint:      rawNode.Data.CSData.LeafCert.Fingerprint,
			NotBefore:        time.Unix(int64(rawNode.Data.CSData.LeafCert.NotBefore), 0).Format(time.RFC3339),
			NotAfter:         time.Unix(int64(rawNode.Data.CSData.LeafCert.NotAfter), 0).Format(time.RFC3339),
			CN:               rawNode.Data.CSData.LeafCert.Subject.CN,
			SourceName:       rawNode.Data.CSData.Source.Name,
			BasicConstraints: rawNode.Data.CSData.LeafCert.Extensions.BasicConstrains,
		}
		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,omitempty"`
	NodeType  string    `json:"nodeType,omitempty"`
	FullPaste FullPaste `json:"fullPaste,omitempty"`
	Created   string    `json:"create,omitempty"`
	Modified  string    `json:"modified,omitempty"`
}

// FullPaste wrapes meta and information from Pastebin.
type FullPaste struct {
	Meta     PasteMeta `json:"meta,omitempty"`
	Full     string    `json:"full,omitempty"`
	NodeType string    `json:"nodeType,omitempty"`
}

// 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(),
		NodeType:  "pastebin",
		FullPaste: *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,omitempty"`
	NodeType string         `json:"nodeType,omitempty"`
	HostData ShodanHostData `json:"hostData,omitempty"`
	Created  string         `json:"created,omitempty"`
	Modified string         `json:"modified,omitempty"`
}

// ShodanHostData is a copy of the structure in the go shodan library. It's a
// workaround to have more control on the data send.
type ShodanHostData struct {
	Product   string   `json:"product,omitempty"`
	Hostnames []string `json:"hostnames,omitempty"`
	Version   string   `json:"version,omitempty"`
	Title     string   `json:"title,omitempty"`
	// SSL          *HostSSL               `json:"ssl"`
	IP           string   `json:"ip_str,omitempty"`
	OS           string   `json:"os,omitempty"`
	Organization string   `json:"org,omitempty"`
	ISP          string   `json:"isp,omitempty"`
	CPE          []string `json:"cpe,omitempty"`
	// Data         string   `json:"data,omitempty"`
	ASN        string   `json:"asn,omitempty"`
	Port       int      `json:"port,omitempty"`
	HTML       string   `json:"html,omitempty"`
	Banner     string   `json:"banner,omitempty"`
	Link       string   `json:"link,omitempty"`
	Transport  string   `json:"transport,omitempty"`
	Domains    []string `json:"domains,omitempty"`
	Timestamp  string   `json:"timestamp,omitempty"`
	DeviceType string   `json:"devicetype,omitempty"`
	// Location     *HostLocation          `json:"location"`
	ShodanData map[string]interface{} `json:"_shodan,omitempty"`
	Opts       map[string]interface{} `json:"opts,omitempty"`
}

// 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(),
		NodeType: "shodan_stream",
		HostData: ShodanHostData{
			Product:      data.Product,
			Hostnames:    data.Hostnames,
			Version:      data.Version.String(),
			Title:        data.Title,
			IP:           data.IP.String(),
			OS:           data.OS,
			Organization: data.Organization,
			ISP:          data.ISP,
			CPE:          data.CPE,
			// Data:         data.Data,
			ASN:        data.ASN,
			Port:       data.Port,
			HTML:       data.HTML,
			Banner:     data.Banner,
			Link:       data.Link,
			Transport:  data.Transport,
			Domains:    data.Domains,
			Timestamp:  data.Timestamp,
			DeviceType: data.DeviceType,
			ShodanData: data.ShodanData,
			Opts:       data.Opts,
		},
		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"`
	NodeType string           `json:"nodeType"`
	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(),
		NodeType: "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)
	}
}