initial vulnbot command

Author: breadchris

lunatrace/bsl/ingest-worker/cmd/vulnbot/main.go

@@ -0,0 +1,86 @@
+package main
+
+import (
+	"net/http"
+	"os"
+	"os/signal"
+
+	"github.com/ajvpot/clifx"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+	"github.com/urfave/cli/v2"
+	"go.uber.org/fx"
+
+	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/config/ingestworker"
+	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/dbfx"
+	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/discordfx"
+	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/graphqlfx"
+	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/ml"
+	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/openaifx"
+	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/pineconefx"
+	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/vulnbot"
+)
+
+type Params struct {
+	fx.In
+
+	VulnBot vulnbot.VulnBot
+}
+
+func NewCommand(p Params) clifx.CommandResult {
+	return clifx.CommandResult{
+		Command: &cli.Command{
+			Name:  "discord",
+			Usage: "",
+			Flags: []cli.Flag{},
+			Action: func(ctx *cli.Context) error {
+				log.Info().Msg("Starting Discord Bot")
+				err := p.VulnBot.Start()
+				if err != nil {
+					return err
+				}
+
+				stop := make(chan os.Signal, 1)
+				signal.Notify(stop, os.Interrupt)
+				<-stop
+				return nil
+			},
+		},
+	}
+}
+
+func main() {
+	// TODO (cthompson) this should be configured with an fx module
+	logLevel := zerolog.InfoLevel
+	if os.Getenv("LOG_LEVEL") == "debug" {
+		logLevel = zerolog.DebugLevel
+	}
+	log.Logger = zerolog.New(os.Stderr).With().Timestamp().Logger().Level(logLevel)
+
+	clifx.Main(
+		// TODO (cthompson) move this into an fx module
+		fx.Supply(http.DefaultClient),
+
+		graphqlfx.Module,
+		dbfx.Module,
+		pineconefx.Module,
+		openaifx.Module,
+		discordfx.Module,
+
+		fx.Invoke(discordfx.RegisterCommands),
+
+		fx.Provide(
+			// TODO (cthompson) make a vulnbot config provider
+			ingestworker.NewConfigProvider,
+			ml.NewService,
+			NewCommand,
+			vulnbot.NewVulnBot,
+		),
+
+		fx.Supply(&clifx.AppConfig{
+			Name:    "vulnbot",
+			Usage:   "LunaTrace Vulnerability Bot",
+			Version: "0.0.1",
+		}),
+	)
+}

lunatrace/bsl/ingest-worker/pkg/discordfx/commands.go

@@ -0,0 +1,126 @@
+package discordfx
+
+import (
+	"context"
+	"time"
+
+	"github.com/bwmarrin/discordgo"
+	"github.com/rs/zerolog/log"
+	"go.uber.org/fx"
+)
+
+type HandlerFunc func(ctx context.Context, s *discordgo.Session, i *discordgo.InteractionCreate)
+
+// ApplicationCommandWithHandler is a top level Command with a Handler function.
+// Subcommands are TODO
+// todo: define our own handler func? want to use channels to post back messages.
+type ApplicationCommandWithHandler struct {
+	Command discordgo.ApplicationCommand
+	Handler HandlerFunc
+	// TODO (cthompson) a bit of a hack right now, these should probably be their own type
+	MessageComponent bool
+	GuildID          string
+}
+
+type RegisterCommandsParams struct {
+	fx.In
+	Config    DiscordConfig
+	Session   *discordgo.Session
+	Commands  []*ApplicationCommandWithHandler `group:"command"`
+	Lifecycle fx.Lifecycle
+}
+
+type interactionHelper struct {
+	i *discordgo.InteractionCreate
+}
+
+type InteractionHelper interface {
+	GetInteraction() *discordgo.InteractionCreate
+	Respond() (<-chan *discordgo.WebhookEdit, error)
+	RespondWebhook() (<-chan *discordgo.WebhookEdit, error)
+}
+
+func NewInteractionHelper(i *discordgo.InteractionCreate) InteractionHelper {
+	//return &interactionHelper{i: i}
+	return nil
+}
+
+func RegisterCommands(p RegisterCommandsParams) error {
+	handlerMap := make(map[string]HandlerFunc)
+	registeredCommands := make([]*discordgo.ApplicationCommand, 0, len(p.Commands))
+
+	ctx, _ := context.WithCancel(context.Background())
+
+	// Set up command
+	for _, commandHandler := range p.Commands {
+		handlerMap[commandHandler.Command.Name] = commandHandler.Handler
+	}
+
+	p.Session.AddHandler(func(s *discordgo.Session, i *discordgo.InteractionCreate) {
+		ctx, cancel := context.WithTimeout(ctx, time.Second*15)
+		defer cancel()
+
+		switch i.Type {
+		case discordgo.InteractionApplicationCommand:
+			log.Info().
+				Str("handlerName", i.ApplicationCommandData().Name).
+				Msg("invoking command handler")
+			if h, ok := handlerMap[i.ApplicationCommandData().Name]; ok {
+				h(ctx, s, i)
+			}
+		case discordgo.InteractionMessageComponent:
+			customID := i.MessageComponentData().CustomID
+			log.Info().
+				Str("customID", customID).
+				Msg("invoking message component handler")
+			if h, ok := handlerMap[customID]; ok {
+				h(ctx, s, i)
+			}
+		}
+	})
+
+	p.Lifecycle.Append(fx.Hook{OnStart: func(ctx context.Context) error {
+		for _, v := range p.Commands {
+			if v.MessageComponent {
+				continue
+			}
+
+			ccmd, err := p.Session.ApplicationCommandCreate(p.Config.ApplicationID, v.GuildID, &v.Command)
+			if err != nil {
+				log.Error().Err(err).Msg("failed to register command")
+				return err
+			}
+			log.Info().
+				Str("name", ccmd.Name).
+				Str("id", ccmd.ID).
+				Msg("registered command")
+
+			registeredCommands = append(registeredCommands, ccmd)
+		}
+		return nil
+	},
+	//OnStop: func(ctx context.Context) error {
+	//	cancel()
+	//	registeredCommands, err := p.Session.ApplicationCommands(p.Config.ApplicationID, "")
+	//	if err != nil {
+	//		log.Error().
+	//			Err(err).
+	//			Msg("could not fetch registered commands")
+	//		return err
+	//	}
+	//	for _, v := range registeredCommands {
+	//		err := p.Session.ApplicationCommandDelete(p.Config.ApplicationID, v.GuildID, v.ID)
+	//		if err != nil {
+	//			log.Error().Err(err).Msg("failed to delete command")
+	//			return err
+	//		}
+	//		log.Debug().
+	//			Str("handlerName", v.Name).
+	//			Str("id", v.ID).
+	//			Msg("deleted command")
+	//	}
+	//	return nil
+	//}
+	})
+	return nil
+}

lunatrace/bsl/ingest-worker/pkg/discordfx/config.go

@@ -0,0 +1,30 @@
+package discordfx
+
+import (
+	"github.com/bwmarrin/discordgo"
+	"github.com/rs/zerolog/log"
+	"go.uber.org/config"
+	"go.uber.org/fx"
+)
+
+const ConfigurationKey = "discord"
+
+type ConfigParams struct {
+	fx.In
+	Config config.Provider
+}
+
+type DiscordConfig struct {
+	ApplicationID string           `yaml:"application_id"`
+	Token         string           `yaml:"token"`
+	Intent        discordgo.Intent `yaml:"intent"`
+}
+
+func NewConfig(p ConfigParams) (cfg DiscordConfig, err error) {
+	err = p.Config.Get(ConfigurationKey).Populate(&cfg)
+	if err != nil {
+		log.Error().Err(err).Msg("failed loading config")
+		return
+	}
+	return
+}

lunatrace/bsl/ingest-worker/pkg/discordfx/module.go

@@ -0,0 +1,7 @@
+package discordfx
+
+import (
+	"go.uber.org/fx"
+)
+
+var Module = fx.Options(fx.Provide(NewConfig, NewDiscordSession))

lunatrace/bsl/ingest-worker/pkg/discordfx/session.go

@@ -0,0 +1,63 @@
+package discordfx
+
+import (
+	"github.com/bwmarrin/discordgo"
+	"github.com/rs/zerolog/log"
+	"go.uber.org/fx"
+)
+
+type NewSessionParams struct {
+	fx.In
+	Config DiscordConfig
+}
+
+func NewDiscordSession(p NewSessionParams) (*discordgo.Session, error) {
+	token := "Bot " + p.Config.Token
+
+	s, err := discordgo.New(token)
+	if err != nil {
+		log.Error().Err(err).Msg("failed to create discord session")
+		return nil, err
+	}
+
+	if p.Config.Intent > 0 {
+		s.Identify.Intents = p.Config.Intent
+	}
+
+	if t, err := ParseToken(token); err != nil {
+		log.Error().Err(err).Msg("failed to parse token")
+	} else {
+		log.Info().Str("url", GenerateOAuthURL(t)).Msg("OAuth URL")
+	}
+
+	return s, nil
+}
+
+func LogSessionEvents(s *discordgo.Session) {
+	s.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) {
+		log.Info().
+			Interface("event", m).
+			Msg("message create")
+	})
+	//d.session.AddHandler(func(s *discordgo.Session, r *discordgo.Ready) {
+	//	log.Info().
+	//		Str("username", s.State.User.Username).
+	//		Str("discriminator", s.State.User.Discriminator).
+	//		Msg("connected")
+	//})
+	//d.session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageEdit) {
+	//	log.Info().
+	//		Interface("event", m).
+	//		Msg("message edit")
+	//})
+	//d.session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageDelete) {
+	//	log.Info().
+	//		Interface("event", m).
+	//		Msg("message delete")
+	//})
+	//d.session.AddHandler(func(s *discordgo.Session, m *discordgo.PresenceUpdate) {
+	//	log.Info().
+	//		Interface("event", m).
+	//		Msg("presence update")
+	//})
+}

lunatrace/bsl/ingest-worker/pkg/discordfx/types.go

@@ -0,0 +1,81 @@
+package discordfx
+
+import (
+	"errors"
+	"net/url"
+	"strconv"
+	"strings"
+)
+
+// Guild represents a Guild ID or Guild URL in the config.
+// TODO Make this work properly. Currently requires the user to truncate the channel if they want URL notation.
+type Guild string
+
+func (g *Guild) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var s string
+	if err := unmarshal(&s); err != nil {
+		return err
+	}
+
+	// If it's numeric, just return it.
+	if _, err := strconv.Atoi(s); err == nil {
+		*g = Guild(s)
+		return nil
+	}
+
+	u, err := url.Parse(s)
+	if err != nil {
+		return err
+	}
+
+	urlParts := strings.Split(u.Path, "/")
+
+	if len(urlParts) < 3 {
+		return errors.New("invalid guild url")
+	}
+
+	*g = Guild(urlParts[2])
+
+	// The result must be numeric.
+	if _, err := strconv.Atoi(s); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Channel represents a Channel ID or Channel URL in the config.
+type Channel string
+
+func (c *Channel) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var s string
+	if err := unmarshal(&s); err != nil {
+		return err
+	}
+
+	// If it's numeric, just return it.
+	if _, err := strconv.Atoi(s); err == nil {
+		*c = Channel(s)
+		return nil
+	}
+
+	u, err := url.Parse(s)
+	if err != nil {
+		return err
+	}
+
+	urlParts := strings.Split(u.Path, "/")
+
+	if len(urlParts) < 4 {
+		return errors.New("invalid channel url")
+	}
+
+	*c = Channel(urlParts[3])
+
+	// The result must be numeric.
+	if _, err := strconv.Atoi(string(*c)); err != nil {
+		return err
+	}
+
+	return nil
+}