<|start_filename|>cmd/sinkholecli/main.go<|end_filename|>
package main

import (
    "fmt"
    "os"
    "flag"
    "time"
    "crypto/rand"

    "github.com/go-chi/jwtauth"

    "github.com/scrapbird/sinkholed/internal/config"
)

func printHelp() {
    fmt.Println("Usage of " + os.Args[0] + ":")
    fmt.Println(os.Args[0] + "<command> [<args>]\n")
    fmt.Printf("Subcommands:\n  gensecret\n  genjwt\n")
}

func generateRandomString(length int) (string, error) {
    const runes = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
    bytes := make([]byte, length)
    _, err := rand.Read(bytes)
    if err != nil {
        return "", err
    }

    for i, b := range bytes {
        bytes[i] = runes[b%byte(len(runes))]
    }

    return string(bytes), nil
}

func gensecret(length int) int {
    secret, err := generateRandomString(length)
    if err != nil {
        fmt.Println("Failed to generate random string", err)
        return 1
    }
    fmt.Println(secret)
    return 0
}

func genjwt(configPath string, expiryDays int) int {
    // Initialize config
    cfg, err := config.InitConfig(configPath, true)
    if err != nil {
        fmt.Println("Configuration error", err)
        return 1
    }

    if cfg.JwtSecret == "" {
        fmt.Println("Please set SINKHOLED_JWTSECRET in your environment")
        return 1
    }

    jti, err := generateRandomString(16)
    if err != nil {
        fmt.Println("Failed to generate random string for jti")
        return 1
    }

    var claims jwtauth.Claims
    if expiryDays == 0 {
        claims = jwtauth.Claims{
            "jti": jti,
        }
    } else {
        claims = jwtauth.Claims{
            "exp": time.Now().AddDate(0, 0, expiryDays).Unix(),
            "jti": jti,
        }
    }

    _, tokenString, _ := cfg.JwtAuth.Encode(claims)
    fmt.Println(tokenString)
    return 0
}

func main() {
    // // Parse command line arguments
    gensecretCommand := flag.NewFlagSet("gensecret", flag.ExitOnError)
    secretSizeFlag := gensecretCommand.Int("length", 64, "Length of the JWT secret")
    genjwtCommand := flag.NewFlagSet("genjwt", flag.ExitOnError)
    expiryTimeFlag := genjwtCommand.Int("length", 0, "Number of days until the JWT expires, 0 for never")
    configFlag := genjwtCommand.String("config", "/etc/sinkholed/sinkholed.yml", "sinkholed config file")

    if len(os.Args) < 2 {
        printHelp()
        os.Exit(1)
    }

    switch os.Args[1] {
    case "-h":
        printHelp()
        os.Exit(0)
    case "gensecret":
        gensecretCommand.Parse(os.Args[2:])
        os.Exit(gensecret(*secretSizeFlag))
    case "genjwt":
        genjwtCommand.Parse(os.Args[2:])
        os.Exit(genjwt(*configFlag, *expiryTimeFlag))
    default:
        printHelp()
        os.Exit(2)
    }
}


<|start_filename|>plugins/httpd/main.go<|end_filename|>
package main

import (
    "fmt"
    "strings"
    "io/ioutil"
    "net/http"
    "net/http/httputil"
    "time"

    log "github.com/sirupsen/logrus"

    "github.com/spf13/viper"

    "github.com/scrapbird/sinkholed/pkg/core"
    "github.com/scrapbird/sinkholed/pkg/plugin"
)

type config struct {
    Timeout int
    LogPath string
    LogLevel string
    ListenAddress string
    TLSListenAddress string
    TLSCertFile string
    TLSKeyFile string
}

type httpdPlugin struct {
    plugin.Plugin
    cfg *config
    downstream chan<- *core.Event
}

func (p *httpdPlugin) wildcardResponse(w http.ResponseWriter, req *http.Request) {
    log.Infoln(req.RequestURI)

    // Note that there are some gotchas here if you require an exact request:
    // DumpRequest returns the given request in its HTTP/1.x wire representation. It should 
    // only be used by servers to debug client requests. The returned representation is an 
    // approximation only; some details of the initial request are lost while parsing it 
    // into an http.Request. In particular, the order and case of header field names are 
    // lost. The order of values in multi-valued headers is kept intact. HTTP/2 requests 
    // are dumped in HTTP/1.x form, not in their original binary representations.

    // If body is true, DumpRequest also returns the body. To do so, it consumes req.Body 
    // and then replaces it with a new io.ReadCloser that yields the same bytes. If 
    // DumpRequest returns an error, the state of req is undefined.
    rawRequest, err := httputil.DumpRequest(req, true)
    if err != nil {
        log.Errorln("Failed to get raw request", err)
    }

    body, err := ioutil.ReadAll(req.Body)
    if err != nil {
        log.Errorln("Failed to read request body", err)
    }

    metadata := map[string]interface{}{
        "SourceIp": req.RemoteAddr,
        "Raw": string(rawRequest),
        "Headers": req.Header,
        "Host": req.Host,
        "Cookies": req.Cookies(),
        "Body": string(body),
        "Method": req.Method,
        "URL": req.URL,
        "UserAgent": req.UserAgent(),
        "Protocol": req.Proto,
    }

    if req.TLS != nil {
        metadata["TLS"] = map[string]interface{}{
            "Version": req.TLS.Version,
            "CipherSuite": req.TLS.CipherSuite,
            "NegotiatedProtocol": req.TLS.NegotiatedProtocol,
        }
    }

    event := &core.Event{
        Type: "request",
        Source: "http",
        Timestamp: time.Now(),
        Metadata: metadata,
    }

    p.downstream <- event

    fmt.Fprintln(w, "Ok")
}

func (p *httpdPlugin) httpdWorker() {
    log.Infoln("Starting HTTP listener on", p.cfg.ListenAddress)
    log.Fatal(http.ListenAndServe(p.cfg.ListenAddress, nil))
}

func (p *httpdPlugin) httpsdWorker() {
    log.Infoln("Starting HTTPS listener on", p.cfg.TLSListenAddress)

    if p.cfg.TLSCertFile == "" {
        log.Errorln("Not listening on HTTPS. TLSCertFile config missing. (Env var: SINKHOLED_HTTPD_TLSCERTFILE)")
        return
    }
    if p.cfg.TLSKeyFile == "" {
        log.Errorln("Not listening on HTTPS. TLSKeyFile config missing. (Env var: SINKHOLED_HTTPD_TLSKEYFILE)")
        return
    }

    log.Fatal(http.ListenAndServeTLS(p.cfg.TLSListenAddress, p.cfg.TLSCertFile, p.cfg.TLSKeyFile, nil))
}

func (p *httpdPlugin) Init(cfg *viper.Viper, downstream chan<- *core.Event) error {
    // Store properties
    p.downstream = downstream

    log.Println("Initializing httpd plugin")

    // Enable env variable overrides
    cfg.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
    cfg.SetEnvPrefix("SINKHOLED_HTTPD")
    cfg.AutomaticEnv()

    // Set default timeout
    cfg.SetDefault("Timeout", 30)

    // Workaround for viper not reading keys from env unless they are manually Get'd
    cfg.Set("Timeout", cfg.Get("Timeout"))
    cfg.Set("ListenAddress", cfg.Get("ListenAddress"))
    cfg.Set("TLSListenAddress", cfg.Get("TLSListenAddress"))
    cfg.Set("TLSCertFile", cfg.Get("TLSCertFile"))
    cfg.Set("TLSKeyFile", cfg.Get("TLSKeyFile"))

    // Parse the config
    var c config
    cfg.Unmarshal(&c)
    p.cfg = &c

    // Start the server
    http.HandleFunc("/", p.wildcardResponse)
    go p.httpdWorker()
    if p.cfg.TLSListenAddress != "" {
        go p.httpsdWorker()
    }

    return nil
}

func (p *httpdPlugin) Halt() error {
    return nil
}

// Export the plugin
var Plugin httpdPlugin



<|start_filename|>plugins/elasticsearch/main.go<|end_filename|>
package main

import (
    "time"
    "sync"
    "strings"
    "encoding/json"
    "net/http"
    "net"
    "crypto/tls"

    log "github.com/sirupsen/logrus"

    "github.com/spf13/viper"

    "github.com/elastic/go-elasticsearch/v8"

    "github.com/scrapbird/sinkholed/pkg/core"
    "github.com/scrapbird/sinkholed/pkg/plugin"
)

const indexMapping = `{
    "mappings" : {
        "sinkholed": {
            "properties": {
                "type" : {
                    "type": "text"
                },
                "timestamp": {
                    "type": "date"
                },
                "samples": {
                    "type": "nested"
                },
                "source": {
                    "type": "text"
                },
                "metadata": {
                    "type": "object",
                    "dynamic": true
                }
            }
        }
    }
}`

type esSample struct {
    Sha256 string `json:"sha256"`
    FileName string `json:"filename"`
}

type esMetadata interface{}

type esEvent struct {
    Type string `json:"type"`
    Timestamp time.Time `json:"timestamp"`
    Samples []esSample `json:"samples"`
    Source string `json:"source"`
    Metadata esMetadata `json:"metadata"`
}

type config struct {
    Addresses []string
    Username string
    Password string
    Aws bool
}

type esPlugin struct {
    plugin.Plugin
    cfg *config
    events chan *core.Event
    es *elasticsearch.Client
    connectMux sync.Mutex
    disconnect chan bool
}

func (p *esPlugin) Init(cfg *viper.Viper, downstream chan<- *core.Event) error {
    log.Println("Initializing elasticsearch plugin")
    p.events = make(chan *core.Event)

    // Enable env variable overrides
    cfg.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
    cfg.SetEnvPrefix("SINKHOLED_ES")
    cfg.AutomaticEnv()

    // Workaround for viper not reading keys from env unless they are manually Get'd
    cfg.Set("Addresses", cfg.Get("Addresses"))
    cfg.Set("Password", cfg.Get("Password"))
    cfg.Set("Username", cfg.Get("Username"))
    cfg.Set("Aws", cfg.Get("Aws"))

    // Parse the config
    var c config
    cfg.Unmarshal(&c)
    p.cfg = &c

    esConfig := elasticsearch.Config{
        Addresses: p.cfg.Addresses,
        Password: <PASSWORD>,
        Username: p.cfg.Username,
    }

    // Support AWS hosted elasticsearch
    if p.cfg.Aws {
        log.Info("AWS hosted elasticsearch detected, accepting self signed certs")
        esConfig.Transport = &http.Transport{
            MaxIdleConnsPerHost:   10,
            ResponseHeaderTimeout: time.Second,
            DialContext:           (&net.Dialer{
                Timeout: 30 * time.Second,
                KeepAlive: 30 * time.Second,
            }).DialContext,
            TLSClientConfig: &tls.Config{
                InsecureSkipVerify: true,
            },
        }
    }

    p.connectMux.Lock()
    go func() {
        defer p.connectMux.Unlock()
        // Connect to elasticsearch
        for {
            log.Println("Attempting to connect to elasticsearch")
            es, err := elasticsearch.NewClient(esConfig)
            if err == nil {
                p.es = es
                log.Println("Connected to elasticsearch")
                break
            }
            log.Errorln("Failed to connect, retrying in 5 seconds")
            time.Sleep(5 * time.Second)
        }

        var indexDate string
        var indexName string

        for {
            select {
            case event := <-p.events:
                // Check if the date has changed
                currentDate := time.Now().Format("2006-01-02")
                if indexDate != currentDate {
                    indexDate = currentDate
                    indexName = "sinkholed-" + indexDate

                    // Check if index and mappings exist - if not create them
                    resp, err := p.es.Indices.Exists([]string{indexName})
                    if err != nil {
                        log.Errorln("Error checking if index exists", err)
                    } else {
                        log.Println(resp)
                        if resp.StatusCode == 404 {
                            // Check if date has changed after the select statement unblocks
                            log.Println("Index doesn't already exist, creating...")
                            res, err := p.es.Indices.Create(indexName, p.es.Indices.Create.WithBody(strings.NewReader(indexMapping)))
                            if err != nil {
                                log.Errorln("Failed to create index", err)
                            } else {
                                log.Println("Result: ", res)
                            }
                        } else {
                            log.Println("Index already exists, res: ", resp.StatusCode)
                        }
                    }

                }

                // Mutate the event for es
                var esEvent esEvent
                for i := range event.Samples {
                    var sample esSample
                    sample.Sha256 = event.Samples[i].Sha256
                    sample.FileName = event.Samples[i].FileName
                    esEvent.Samples = append(esEvent.Samples, sample)
                }
                esEvent.Type = event.Type
                esEvent.Timestamp = event.Timestamp
                esEvent.Source = event.Source
                esEvent.Metadata = event.Metadata

                body, err := json.Marshal(&esEvent)
                if err != nil {
                    log.Errorln("Error marshaling event:", event)
                    continue
                }

                _, err = p.es.Index(indexName, strings.NewReader(string(body)))
                if err != nil {
                    log.Errorln("Error indexing event:", err)
                }

            case <-p.disconnect:
                break
            }

        }
    }()
    return nil
}

func (p *esPlugin) Halt() error {
    p.disconnect <- true
    return nil
}

func (p *esPlugin) Inbox(event *core.Event) error {
    p.events <- event
    return nil
}

// Export the plugin
var Plugin esPlugin



<|start_filename|>plugins/samplefs/main.go<|end_filename|>
package main

import (
    "os"
    "io/ioutil"
    "path/filepath"

    "github.com/spf13/viper"

    log "github.com/sirupsen/logrus"

    "github.com/scrapbird/sinkholed/pkg/core"
    "github.com/scrapbird/sinkholed/pkg/plugin"
)

type config struct {
    DestDir string
}

type sampleFs struct {
    plugin.Plugin
    cfg *config
    x int
}

func (p *sampleFs) Init(cfg *viper.Viper, downstream chan<- *core.Event) error {
    log.Println("Initializing sampleFs plugin")

    // Parse the config
    var c config
    cfg.Unmarshal(&c)
    p.cfg = &c

    return nil
}

func (p *sampleFs) Halt() error {
    return nil
}


func (p *sampleFs) Inbox(event *core.Event) error {
    for _, sample := range event.Samples {
        filePath := filepath.Join(p.cfg.DestDir, sample.Sha256)

        if _, err := os.Stat(filePath); os.IsNotExist(err) {
            err := ioutil.WriteFile(filePath, sample.Data, 0644)
            if err != nil {
                log.Errorln("Failed to write file", filePath)
            } else {
                log.Println("Wrote file", filePath)
            }
        }

    }
    return nil
}

// Export the plugin
var Plugin sampleFs



<|start_filename|>plugins/downstream/main.go<|end_filename|>
package main

// Sinkholed Downstream Plugin

import (
    "time"
    "bytes"
    "net/http"
    "strings"
    "encoding/json"

    "github.com/spf13/viper"

    log "github.com/sirupsen/logrus"

    "github.com/scrapbird/sinkholed/pkg/core"
    "github.com/scrapbird/sinkholed/pkg/plugin"
)

type config struct {
    DestAddress string
    Jwt string
    TimeoutSeconds time.Duration
}

type downstreamSinkholed struct {
    plugin.Plugin
    cfg *config
}

func (p *downstreamSinkholed) Init(cfg *viper.Viper, downstream chan<- *core.Event) error {
    log.Println("Initializing downstreamSinkholed plugin")

    // Enable env variable overrides
    cfg.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
    cfg.SetEnvPrefix("SINKHOLED_DOWNSTREAM")
    cfg.AutomaticEnv()

    // Set default timeout seconds
    cfg.SetDefault("TimeoutSeconds", 30)

    // Workaround for viper not reading keys from env unless they are manually Get'd
    cfg.Set("DestAddress", cfg.Get("DestAddress"))
    cfg.Set("Jwt", cfg.Get("Jwt"))
    cfg.Set("TimeoutSeconds", cfg.Get("TimeoutSeconds"))

    // Parse the config
    var c config
    cfg.Unmarshal(&c)
    p.cfg = &c

    return nil
}

func (p *downstreamSinkholed) Halt() error {
    return nil
}

func (p *downstreamSinkholed) Inbox(event *core.Event) error {
    // Pipe event to downstream sinkholed instance
    j, err := json.Marshal(event)

    req, err := http.NewRequest("POST", p.cfg.DestAddress + "/api/v1/event", bytes.NewBuffer(j))
    req.Header.Add("Authorization", "Bearer " + p.cfg.Jwt)

    client := &http.Client{
        Timeout: time.Second * p.cfg.TimeoutSeconds,
    }
    resp, err := client.Do(req)
    if err != nil {
        log.Errorln("Failed to send event to downstream:", err)
        return nil
    }
    defer resp.Body.Close()

    if resp.StatusCode != 200 {
        log.Errorln("Failed to send event to downstream. Got http status:", resp.Status)
    }

    return nil
}

// Export the plugin
var Plugin downstreamSinkholed



<|start_filename|>plugins/smtpd/main.go<|end_filename|>
package main

import (
    "time"
    "bytes"
    "strings"

    log "github.com/sirupsen/logrus"

    "github.com/scrapbird/go-guerrilla"
    "github.com/scrapbird/go-guerrilla/backends"
    "github.com/scrapbird/go-guerrilla/mail"

    "github.com/spf13/viper"

    "github.com/scrapbird/sinkholed/pkg/core"
    "github.com/scrapbird/sinkholed/pkg/plugin"

    "github.com/jhillyerd/enmime"
)

type config struct {
    Timeout int
    LogPath string
    LogLevel string
    ListenAddress string
    HostMetadata map[string]map[string]interface{}
    HostTags map[string][]string
}

type smtpdPlugin struct {
    plugin.Plugin
    cfg *config
    downstream chan<- *core.Event
}

func (p *smtpdPlugin) smtpdWorker() {
    cfg := &guerrilla.AppConfig{
        LogFile: p.cfg.LogPath,
        LogLevel: p.cfg.LogLevel,
        AllowedHosts: []string{"."},
    }

    sc := guerrilla.ServerConfig{
        ListenInterface: p.cfg.ListenAddress,
        IsEnabled: true,
        Timeout: p.cfg.Timeout,
    }
    cfg.Servers = append(cfg.Servers, sc)
    bcfg := backends.BackendConfig{
        "save_workers_size":  3,
        "save_process":      "HeadersParser|Header|Hasher|Sinkholed",
        "primary_mail_host" : "smtp.google.com",
    }
    cfg.BackendConfig = bcfg

    d := guerrilla.Daemon{Config: cfg}

    var sinkholedProcessor = func() backends.Decorator {
        return func(processor backends.Processor) backends.Processor {
            return backends.ProcessWith(
                func(e *mail.Envelope, task backends.SelectTask) (backends.Result, error) {
                    if task == backends.TaskValidateRcpt {
                            return processor.Process(e, task)
                        } else if task == backends.TaskSaveMail {
                            // Annoying that we have to parse the MIME message again
                            // but I couldn't find an easy way using only go-guerrilla
                            envelope, err := enmime.ReadEnvelope(bytes.NewReader(e.Data.Bytes()))
                            if err != nil {
                                log.Errorln("Failed to parse envelope:", err)
                                return processor.Process(e, task)
                            }

                            samples := []core.Sample{}
                            for _, attachment := range envelope.Attachments {
                                sample := core.NewSample(attachment.FileName, attachment.Content)
                                samples = append(samples, *sample)
                            }

                            var recipients []string
                            for _, r := range e.RcptTo {
                                recipients = append(recipients, r.String())
                            }

                            metadata := map[string]interface{}{
                                "SourceIp": e.RemoteIP,
                                "MailFrom": e.MailFrom.String(),
                                "Data": e.Data.String(),
                                "Body": envelope.Text,
                                "Subject": e.Subject,
                                "RcptTo": recipients,
                            }

                            // Set host tags if set in config file
                            if tags, ok := p.cfg.HostTags[e.RemoteIP]; ok {
                                metadata["tags"] = tags
                            }

                            // Override metadata if set in config file
                            if configMetadata, ok := p.cfg.HostMetadata[e.RemoteIP]; ok {
                                log.Infoln("Found config metadata")
                                for k, v := range configMetadata {
                                    metadata[k] = v
                                }
                            }

                            event := &core.Event{
                                Type: "email",
                                Source: "smtp",
                                Timestamp: time.Now(),
                                Metadata: metadata,
                                Samples: samples,
                            }

                            p.downstream <-event
                            // call the next processor in the chain
                            return processor.Process(e, task)
                        }
                        return processor.Process(e, task)
                    },
                )
            }
        }
    d.AddProcessor("Sinkholed", sinkholedProcessor)

    log.Infoln("Starting smtp server")
    err := d.Start()

    if err != nil {
        log.Error("Failed to start SMTP server", err)
    }
}

func (p *smtpdPlugin) Init(cfg *viper.Viper, downstream chan<- *core.Event) error {
    // Store properties
    p.downstream = downstream

    log.Println("Initializing smtpd plugin")

    // Enable env variable overrides
    cfg.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
    cfg.SetEnvPrefix("SINKHOLED_ES")
    cfg.AutomaticEnv()

    // Set default timeout
    cfg.SetDefault("Timeout", 30)

    // Workaround for viper not reading keys from env unless they are manually Get'd
    cfg.Set("Timeout", cfg.Get("Timeout"))
    cfg.Set("ListenAddress", cfg.Get("ListenAddress"))
    cfg.Set("HostMetadata", cfg.Get("HostMetadata"))
    cfg.Set("HostTags", cfg.Get("HostTags"))

    // Parse the config
    var c config
    cfg.Unmarshal(&c)
    p.cfg = &c
    // Workaround for https://github.com/spf13/viper/issues/324
    cfg.UnmarshalKey("HostTags", &p.cfg.HostTags)
    cfg.UnmarshalKey("HostMetadata", &p.cfg.HostMetadata)

    // Start the server
    go p.smtpdWorker()

    return nil
}

func (p *smtpdPlugin) Halt() error {
    return nil
}

// Export the plugin
var Plugin smtpdPlugin



<|start_filename|>internal/api/v1/routes.go<|end_filename|>
package v1

import (
    "fmt"
    "net/http"
    "encoding/json"
    "crypto/sha256"

    "github.com/go-chi/chi"
    "github.com/go-chi/render"
    "github.com/go-chi/jwtauth"

    // log "github.com/sirupsen/logrus"

    "github.com/scrapbird/sinkholed/internal/config"
    "github.com/scrapbird/sinkholed/pkg/plugin"
    "github.com/scrapbird/sinkholed/pkg/core"
)

type APIResponse struct {
    Message string `json:"message"`
}

func Routes(cfg *config.Config, pluginManager *plugin.PluginManager) *chi.Mux {
    router := chi.NewRouter()

    // Public routes
    router.Group(func (r chi.Router) {
        r.Get("/healthcheck", Healthcheck)
    })

    // Private routes
    router.Group(func (r chi.Router) {
        r.Use(jwtauth.Verifier(cfg.JwtAuth))
        r.Use(jwtauth.Authenticator)
        r.Post("/event", PostEvent(cfg, pluginManager))
    })

    return router
}

func Healthcheck(w http.ResponseWriter, r *http.Request) {
    resp := APIResponse{
        Message: "ok",
    }
    render.JSON(w, r, resp)
}

func PostEvent(cfg *config.Config, pluginManager *plugin.PluginManager) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        if r.Body == nil {
            http.Error(w, "Please send a request body", 400)
            return
        }

        var event core.Event

        err := json.NewDecoder(r.Body).Decode(&event)
        if err != nil {
            http.Error(w, err.Error(), 400)
            return
        }

        for i := range event.Samples {
            // Calculate hash
            sha256sum := sha256.Sum256(event.Samples[i].Data)
            event.Samples[i].Sha256 = fmt.Sprintf("%x", sha256sum)
        }

        pluginManager.EmitEvent(&event)

        resp := APIResponse{
            Message: "ok",
        }
        render.JSON(w, r, resp)
    }
}



<|start_filename|>pkg/core/core.go<|end_filename|>
package core

import (
    "fmt"
    "time"
    "crypto/sha256"
)

// The type used to represent a binary sample file.
type Sample struct {
    // Binary representation of the sample file
    Data []byte `json:"data"`
    // Sha256 of the Data
    Sha256 string `json:"sha256"`
    // Original file name
    FileName string `json:"filename"`
}

// Metadata is defined as an interface{} so that we can use any method for attaching 
// metadata to an event. For example, an array of string tags or a more complicated map
// of data.
type Metadata interface{}

// The type that is used to represent events passed through the sinkhole.
type Event struct {
    // This string should indicate the type of event, 'email', 'dropped file' etc.
    Type string `json:"type"`
    // A time stamp of the time the event happened
    Timestamp time.Time `json:"timestamp"`
    // An array of samples, (attachments, the file dropped etc).
    Samples []Sample `json:"samples"`
    // The source of the event, (name of the plugin that generated the event, unenforced)
    Source string `json:"source"`
    // Any metadata about the event, an array of tags or a more complicated map detailing 
    // the botnet it came from etc etc
    Metadata Metadata `json:"metadata"`
}

// Used to create a new sample representation
//
// This function will automatically calculate the sha256 sum of the sample
func NewSample(filename string, data []byte) *Sample {
    sample := Sample{
        FileName: filename,
        Data: data,
    }
    sha256sum := sha256.Sum256(data)
    sample.Sha256 = fmt.Sprintf("%x", sha256sum)
    return &sample
}



<|start_filename|>plugins/dns/main.go<|end_filename|>
package main

import (
    "fmt"
    "strings"
    "time"

    log "github.com/sirupsen/logrus"

    "github.com/miekg/dns"

    "github.com/spf13/viper"

    "github.com/scrapbird/sinkholed/pkg/core"
    "github.com/scrapbird/sinkholed/pkg/plugin"
)

type Record struct {
    Domain string
    Address string
}

type config struct {
    Timeout int
    LogPath string
    LogLevel string
    ListenAddress string
    Records []Record
}

type dnsPlugin struct {
    plugin.Plugin
    cfg *config
    downstream chan<- *core.Event
    server *dns.Server
    records map[string]string
}

func (p *dnsPlugin) parseQuery(m *dns.Msg) {
    for _, q := range m.Question {
        switch q.Qtype {
        case dns.TypeA:
            log.Infoln("Answering DNS query for", q.Name)
            ip := p.records[q.Name]

            if ip != "" {
                rr, err := dns.NewRR(fmt.Sprintf("%s A %s", q.Name, ip))
                if err == nil {
                    m.Answer = append(m.Answer, rr)
                }
            }
        }
    }
}

func (p *dnsPlugin) handleDnsRequest(w dns.ResponseWriter, r *dns.Msg) {
    m := new(dns.Msg)
    m.SetReply(r)
    m.Compress = false

    metadata := map[string]interface{}{
        "Raw": r.String(),
    }

    event := &core.Event{
        Type: "request",
        Source: "dns",
        Timestamp: time.Now(),
        Metadata: metadata,
    }

    p.downstream <- event

    switch r.Opcode {
    case dns.OpcodeQuery:
        p.parseQuery(m)
    }

    w.WriteMsg(m)
}

func (p *dnsPlugin) Init(cfg *viper.Viper, downstream chan<- *core.Event) error {
    // Store properties
    p.downstream = downstream

    log.Println("Initializing dns plugin")

    // Enable env variable overrides
    cfg.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
    cfg.SetEnvPrefix("SINKHOLED_DNS")
    cfg.AutomaticEnv()

    // Set default timeout
    cfg.SetDefault("Timeout", 30)

    // Workaround for viper not reading keys from env unless they are manually Get'd
    cfg.Set("Timeout", cfg.Get("Timeout"))
    cfg.Set("ListenAddress", cfg.Get("ListenAddress"))
    cfg.Set("Records", cfg.Get("Records"))

    // Parse the config
    var c config
    cfg.Unmarshal(&c)
    p.cfg = &c

    // This is a workaround because viper can't contain dots in config key names (domains)
    // Process records into a map
    p.records = map[string]string{}
    for _, record := range p.cfg.Records {
        p.records[record.Domain] = record.Address
    }

    // Attach request handler func
    dns.HandleFunc(".", p.handleDnsRequest)

    // Start server
    p.server = &dns.Server{
        Addr: p.cfg.ListenAddress,
        Net: "udp",
        ReadTimeout: time.Second * time.Duration(p.cfg.Timeout),
        WriteTimeout: time.Second * time.Duration(p.cfg.Timeout),
    }
    log.Infoln("Starting DNS listener on", p.cfg.ListenAddress)

    go func() {
        err := p.server.ListenAndServe()
        defer p.server.Shutdown()
        if err != nil {
            log.Fatalf("Failed to start DNS server: %s\n ", err.Error())
        }
    }()

    return nil
}

func (p *dnsPlugin) Halt() error {
    p.server.Shutdown()
    return nil
}

// Export the plugin
var Plugin dnsPlugin




<|start_filename|>Dockerfile<|end_filename|>
FROM golang:1.11.6-alpine as build

ENV GOPATH=""

RUN apk add --update --no-cache git build-base bash

RUN mkdir /opt/sinkholed

WORKDIR /opt/sinkholed

COPY ./cmd cmd
COPY ./config config
COPY ./internal internal
COPY ./pkg pkg
COPY ./plugins plugins
COPY ./build.sh build.sh
COPY ./go.mod go.mod
COPY ./go.sum go.sum

RUN ./build.sh

FROM alpine

RUN apk add --update --no-cache libcap

# Add sinkholed user and group
RUN addgroup -S sinkholed \
    && adduser -S sinkholed -G sinkholed

RUN mkdir -p /var/lib/sinkholed/samples
VOLUME /var/lib/sinkholed/samples

VOLUME /etc/sinkholed/certs

RUN mkdir -p /usr/local/lib/sinkholed
COPY --from=build /opt/sinkholed/bin/*.so /usr/local/lib/sinkholed/

COPY --from=build /opt/sinkholed/bin/sinkholed /usr/local/bin/sinkholed

COPY ./config/sinkholed.yml /etc/sinkholed/sinkholed.yml
COPY ./entrypoint.sh /usr/local/bin/entrypoint.sh

RUN setcap CAP_NET_BIND_SERVICE=+eip /usr/local/bin/sinkholed

# Create log dir and set permissions
RUN mkdir /var/log/sinkholed \
    && chown sinkholed:sinkholed /var/log/sinkholed

USER sinkholed

ENTRYPOINT /usr/local/bin/entrypoint.sh /usr/local/bin/sinkholed /var/log/sinkholed/sinkholed.log



<|start_filename|>pkg/plugin/plugin.go<|end_filename|>
// This package contains the implementation of the plugin system for extending sinkholed.
//
// A sinkholed plugin is just a go plugin. A go plugin should reside in the main package
// and implement at least the Plugin interface as defined in this package.
//
// There are three kinds of plugins in sinkholed:
//
// Downstream plugins:
// Plugins that recieve events. These plugins could be used to broadcast events to other
// services, store events etc.
//
// Upstream plugins:
// Plugins that produce events to be consumed by downstream plugins. These plugins make use
// of the event channel passed into their Init method, piping all events they produce into
// it. This will cause the event to be consumed and processed by sinkholed.
//
// Middleware plugins:
// Plugins that filter or change events. These act like middleware and can be used to drop
// events, causing them to not be processed or modify events to do actions such as adding
// extra metadata or removing unwanted data.
//
// Any plugin can implment more than one or all of the previously described interfaces and
// sinkholed will automatically figure out what capability the plugin has. This means that a
// single plugin can act as an upstream, downstream and milddleware plugin all at once.
//
// To load a plugin, define a configuration object in the "Plugins" map in the sinkholed
// config file, using the plugin file name without the extension plugin as the key. Example:
//
//  Plugins:
//    "pluginName":
//      ConfigVariable: config value
//
// The configuration object (stored in the main sinkholed config file in the "Plugins" map)
// will be passed to the plugin Init function as a *viper.Viper config object. Included in
// this map will also be two extra values: `LogPath`, which is the path to be used for the 
// log file, and `LogLevel`, which is the log level to use. This can be overridden if you 
// provide a value for LogPath in your plugin's config.
//
// An example plugin is shown below:
//
//  package main
//
//  import (
//      log "github.com/sirupsen/logrus"
//
//      "github.com/spf13/viper"
//  
//      "github.com/scrapbird/sinkholed/pkg/core"
//      "github.com/scrapbird/sinkholed/pkg/plugin"
//  )
//  
//  type demo struct {
//      plugin.Plugin
//      cfg *viper.Viper
//      downstream chan *core.Event
//  }
//  
//  func (p *demo) Init(cfg *viper.Viper, downstream chan *core.Event) error {
//      log.Println("Initializing demo plugin")
//      p.downstream = downstream
//      return nil
//  }
//  
//  func (p *demo) Halt() error {
//      return nil
//  }
//
//  func (p *demo) Inbox(event *core.Event) error {
//      log.Println("New event:", event)
//      return nil
//  }
//  
//  // Export the plugin
//  var Plugin demo
//
// To compile this plugin:
//
//  go build -buildmode=plugin -o main.so main.go
package plugin

import (
    "plugin"
    "sync"
    log "github.com/sirupsen/logrus"

    "github.com/spf13/viper"

    "github.com/scrapbird/sinkholed/internal/config"
    "github.com/scrapbird/sinkholed/pkg/core"
)

// Plugin is the interface that every plugin must implement.
type Plugin interface {
    // Init takes two parameters, the first is a *viper.Viper config object
    // seeded from the plugin config object defined in the sinkholed config
    // file. The second parameter is a channel used to send events downstream
    // and into the sinkhole.
    Init(*viper.Viper, chan<- *core.Event) error
    // Halt should be implemented by every plugin even if it isn't required. This is
    // where the plugin should clean up any used resources (close DB connections etc).
    Halt() error
}

// UpstreamPlugin is the interface that every upstream plugin 
// (plugin that produces events) must implement.
type UpstreamPlugin interface {
    Plugin
}

// DownstreamPlugin is the interface that every downstream plugin 
// (plugin that consumes events) must implement.
type DownstreamPlugin interface {
    Plugin
    // Inbox is used to pass events to downstream plugins. This should be implemented
    // by all plugins wishing to act as a downstream plugin.
    Inbox(event *core.Event) error
}

// MiddlewarePlugin is the interface that every middleware plugin 
// (plugin that filters or changes events before being sent downstream)
// must implement.
type MiddlewarePlugin interface {
    Plugin
    FilterEvent(event *core.Event) (*core.Event, error)
}

// PluginManager manages plugins and handles the flow of events through
// them.
type PluginManager struct {
    cfg          *config.Config
    plugins      []Plugin
    eventChannel chan *core.Event
    running bool
}

// Creates a new plugin manager.
func NewPluginManager(cfg *config.Config) *PluginManager {
    pm := PluginManager{
        cfg:     cfg,
        plugins: make([]Plugin, 0),
        eventChannel: make(chan *core.Event),
    }
    return &pm
}

// Initialize the plugin manager. This starts processing incoming events.
func (pm *PluginManager) Init() {
    go func() {
        for event := range pm.eventChannel {
            log.Debugln("Emitting event:", event)
            pm.EmitEvent(event)
        }
        log.Println("Event channel closed.")
    }()
}

// Loads a plugin from a *.so file.
func (pm *PluginManager) LoadPlugin(pluginPath string, pluginConfig *viper.Viper) {
    log.Println("Loading plugin:", pluginPath)
    p, err := plugin.Open(pluginPath)
    if err != nil {
        log.Errorln("Couldn't open plugin " + pluginPath, err)
        return
    }

    // Lookup plugin exports
    symPlugin, err := p.Lookup("Plugin")
    if err != nil {
        log.Errorln("Couldn't load plugin exports:", err)
        return
    }

    // Assert plugin type
    var plugin Plugin
    plugin, ok := symPlugin.(Plugin)
    if !ok {
        log.Errorln("Invalid plugin type:", pluginPath)
        return
    }

    // Initialize plugin
    err = plugin.Init(pluginConfig, pm.eventChannel)
    if err != nil {
        log.Errorf("Failed to initialize plugin %s: %s\n", pluginPath, err)
        return
    }

    pm.plugins = append(pm.plugins, plugin)
}

// Halts all loaded plugins.
func (pm *PluginManager) HaltAllPlugins() {
    var wg sync.WaitGroup
    wg.Add(len(pm.plugins))
    for _, plugin := range pm.plugins {
        go func(plugin Plugin) {
            defer wg.Done()
            plugin.Halt()
        }(plugin)
    }
    wg.Wait()
}

// Sends an event through all middleware plugins for filtering / modification
func (pm *PluginManager) filterEvent(event *core.Event) *core.Event {
    for _, plugin := range pm.plugins {
        p, ok := plugin.(MiddlewarePlugin)
        if ok {
            // This is a middleware plugin.
            var err error
            if event, err = p.FilterEvent(event); err != nil {
                log.Errorln("Error during middleware filtering in plugin", err)
            }
        }
    }
    return event
}

// Sends an event to all downstream plugins.
func (pm *PluginManager) EmitEvent(event *core.Event) {
    event = pm.filterEvent(event)
    if event != nil {
        for _, plugin := range pm.plugins {
            go func(plugin Plugin) {
                // Assert plugin type.
                p, ok := plugin.(DownstreamPlugin)
                if ok {
                    // This is a downstream plugin.
                    p.Inbox(event)
                }
            }(plugin)
        }
    }
}


<|start_filename|>cmd/sinkholed/main.go<|end_filename|>
package main

import (
    "flag"
    "net/http"
    "os"
    "path/filepath"

    log "github.com/sirupsen/logrus"

    "github.com/go-chi/chi"
    "github.com/go-chi/chi/middleware"
    "github.com/go-chi/render"

    "github.com/scrapbird/sinkholed/internal/api/v1"
    "github.com/scrapbird/sinkholed/internal/config"
    "github.com/scrapbird/sinkholed/pkg/plugin"
)

func Routes(cfg *config.Config, pluginManager *plugin.PluginManager, logger *log.Logger) *chi.Mux {
    var LoggerMiddleware = func(next http.Handler) http.Handler {
        formatter := middleware.DefaultLogFormatter{Logger: logger, NoColor: true}
        return middleware.RequestLogger(&formatter)(next)
    }

    router := chi.NewRouter()
    router.Use(
        render.SetContentType(render.ContentTypeJSON),
        LoggerMiddleware,
        middleware.DefaultCompress,
        middleware.RedirectSlashes,
        middleware.Recoverer,
    )

    // Healthcheck route
    type APIResponse struct {
        Message string `json:"message"`
    }
    router.Get("/", func (w http.ResponseWriter, r *http.Request) {
        resp := APIResponse{
            Message: "ok",
        }
        render.JSON(w, r, resp)
    })

    router.Route("/api", func(r chi.Router) {
        r.Mount("/v1", v1.Routes(cfg, pluginManager))
    })

    return router
}

func main() {
    log.SetFormatter(&log.JSONFormatter{})

    // Parse command line arguments
    var configPath string
    flag.StringVar(&configPath, "config", "/etc/sinkholed/sinkholed.yml", "Config file")
    flag.Parse()

    // Initialize config
    cfg, err := config.InitConfig(configPath)
    if err != nil {
        log.Panicln("Configuration error", err)
    }

    if cfg.JwtSecret == "" {
        log.Fatalln("Please set SINKHOLED_JWTSECRET in your environment")
    }

    // Initialize logger
    lf, err := os.OpenFile(cfg.LogPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0660)
    if err != nil {
        log.Panicln("Failed to open log file:", err)
    }
    log.SetOutput(lf)

    // Set log level
    switch cfg.LogLevel {
    case "debug":
        log.SetLevel(log.DebugLevel)
    case "error":
        log.SetLevel(log.ErrorLevel)
    case "info":
    default:
        log.SetLevel(log.InfoLevel)
    }

    // Initialize plugin manager and plugins
    pluginManager := plugin.NewPluginManager(cfg)
    pluginManager.Init()

    for pluginName, pluginCfg := range cfg.PluginConfigs {
        // Inject log path and log level if they haven't already been set in plugin config
        if pluginCfg.GetString("LogPath") == "" {
            pluginCfg.Set("LogPath", cfg.LogPath)
        }
        if pluginCfg.GetString("LogLevel") == "" {
            pluginCfg.Set("LogLevel", cfg.LogLevel)
        }

        pluginManager.LoadPlugin(filepath.Join(cfg.PluginsPath, pluginName + ".so"), pluginCfg)
    }

    // Initialize router and start API server
    logger := log.New()
    logger.SetOutput(lf)
    router := Routes(cfg, pluginManager, logger)
    log.Fatal(http.ListenAndServe(cfg.ListenAddr, router))
}



<|start_filename|>internal/config/config.go<|end_filename|>
package config

import (
    "fmt"
    "strings"

    "github.com/go-chi/jwtauth"
    "github.com/spf13/viper"
)

type Plugin struct{
    Config *viper.Viper
}

type Constants struct {
    ListenAddr  string
    JwtSecret   string
    LogPath     string
    LogLevel    string
    PluginsPath string
}

type Config struct {
    Constants

    JwtAuth *jwtauth.JWTAuth
    PluginConfigs map[string]*viper.Viper
}

func InitConfig(configPath string, onlyEnvs ...bool) (*Config, error) {
    readConfig := true

    if (len(onlyEnvs) > 0 && onlyEnvs[0]) {
        readConfig = false
    }

    v := viper.New()

    v.SetEnvPrefix("sinkholed")
    v.AutomaticEnv()

    if readConfig {
        v.SetConfigFile(configPath)

        err := v.ReadInConfig()
        if err != nil {
            return &Config{}, err
        }
    }

    viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
    v.SetEnvPrefix("sinkholed")
    v.AutomaticEnv()

    // Workaround for viper not reading kets from env unless they are manually Get'd
    for _, key := range v.AllKeys() {
        val := v.Get(key)
        v.Set(key, val)
    }
    v.Set("ListenAddr", v.Get("ListenAddr"))
    v.Set("JwtSecret", v.Get("JwtSecret"))
    v.Set("LogPath", v.Get("LogPath"))
    v.Set("LogLevel", v.Get("LogLevel"))
    v.Set("PluginsPath", v.Get("PluginsPath"))

    v.SetDefault("ListenAddr", "127.0.0.1:8080")
    v.SetDefault("LogPath", "/var/log/sinkholed.log")

    var cfg Config
    err := v.Unmarshal(&cfg.Constants)

    if readConfig {
        // Iterate over each plugin config and get a viper instance to be passed to the plugin
        cfg.PluginConfigs = make(map[string]*viper.Viper)
        rawPluginConfigs := v.Get("Plugins").(map[string]interface{})

        for pluginPath := range rawPluginConfigs {
            cfg.PluginConfigs[pluginPath] = v.Sub(fmt.Sprintf("Plugins.%s", pluginPath))
        }
    }

    cfg.JwtAuth = jwtauth.New("HS256", []byte(cfg.JwtSecret), nil)

    return &cfg, err
}
