Skip to content

Commit f529d1f

Browse files
breadchrisbreadchris
committed
initial vulnbot command
1 parent 693cde3 commit f529d1f

File tree

9 files changed

+743
-0
lines changed

9 files changed

+743
-0
lines changed

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

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package main
2+
3+
import (
4+
"net/http"
5+
"os"
6+
"os/signal"
7+
8+
"github.com/ajvpot/clifx"
9+
"github.com/rs/zerolog"
10+
"github.com/rs/zerolog/log"
11+
"github.com/urfave/cli/v2"
12+
"go.uber.org/fx"
13+
14+
"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/config/ingestworker"
15+
"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/dbfx"
16+
"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/discordfx"
17+
"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/graphqlfx"
18+
"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/ml"
19+
"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/openaifx"
20+
"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/pineconefx"
21+
"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/vulnbot"
22+
)
23+
24+
type Params struct {
25+
fx.In
26+
27+
VulnBot vulnbot.VulnBot
28+
}
29+
30+
func NewCommand(p Params) clifx.CommandResult {
31+
return clifx.CommandResult{
32+
Command: &cli.Command{
33+
Name: "discord",
34+
Usage: "",
35+
Flags: []cli.Flag{},
36+
Action: func(ctx *cli.Context) error {
37+
log.Info().Msg("Starting Discord Bot")
38+
err := p.VulnBot.Start()
39+
if err != nil {
40+
return err
41+
}
42+
43+
stop := make(chan os.Signal, 1)
44+
signal.Notify(stop, os.Interrupt)
45+
<-stop
46+
return nil
47+
},
48+
},
49+
}
50+
}
51+
52+
func main() {
53+
// TODO (cthompson) this should be configured with an fx module
54+
logLevel := zerolog.InfoLevel
55+
if os.Getenv("LOG_LEVEL") == "debug" {
56+
logLevel = zerolog.DebugLevel
57+
}
58+
log.Logger = zerolog.New(os.Stderr).With().Timestamp().Logger().Level(logLevel)
59+
60+
clifx.Main(
61+
// TODO (cthompson) move this into an fx module
62+
fx.Supply(http.DefaultClient),
63+
64+
graphqlfx.Module,
65+
dbfx.Module,
66+
pineconefx.Module,
67+
openaifx.Module,
68+
discordfx.Module,
69+
70+
fx.Invoke(discordfx.RegisterCommands),
71+
72+
fx.Provide(
73+
// TODO (cthompson) make a vulnbot config provider
74+
ingestworker.NewConfigProvider,
75+
ml.NewService,
76+
NewCommand,
77+
vulnbot.NewVulnBot,
78+
),
79+
80+
fx.Supply(&clifx.AppConfig{
81+
Name: "vulnbot",
82+
Usage: "LunaTrace Vulnerability Bot",
83+
Version: "0.0.1",
84+
}),
85+
)
86+
}

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

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
package discordfx
2+
3+
import (
4+
"context"
5+
"time"
6+
7+
"github.com/bwmarrin/discordgo"
8+
"github.com/rs/zerolog/log"
9+
"go.uber.org/fx"
10+
)
11+
12+
type HandlerFunc func(ctx context.Context, s *discordgo.Session, i *discordgo.InteractionCreate)
13+
14+
// ApplicationCommandWithHandler is a top level Command with a Handler function.
15+
// Subcommands are TODO
16+
// todo: define our own handler func? want to use channels to post back messages.
17+
type ApplicationCommandWithHandler struct {
18+
Command discordgo.ApplicationCommand
19+
Handler HandlerFunc
20+
// TODO (cthompson) a bit of a hack right now, these should probably be their own type
21+
MessageComponent bool
22+
GuildID string
23+
}
24+
25+
type RegisterCommandsParams struct {
26+
fx.In
27+
Config DiscordConfig
28+
Session *discordgo.Session
29+
Commands []*ApplicationCommandWithHandler `group:"command"`
30+
Lifecycle fx.Lifecycle
31+
}
32+
33+
type interactionHelper struct {
34+
i *discordgo.InteractionCreate
35+
}
36+
37+
type InteractionHelper interface {
38+
GetInteraction() *discordgo.InteractionCreate
39+
Respond() (<-chan *discordgo.WebhookEdit, error)
40+
RespondWebhook() (<-chan *discordgo.WebhookEdit, error)
41+
}
42+
43+
func NewInteractionHelper(i *discordgo.InteractionCreate) InteractionHelper {
44+
//return &interactionHelper{i: i}
45+
return nil
46+
}
47+
48+
func RegisterCommands(p RegisterCommandsParams) error {
49+
handlerMap := make(map[string]HandlerFunc)
50+
registeredCommands := make([]*discordgo.ApplicationCommand, 0, len(p.Commands))
51+
52+
ctx, _ := context.WithCancel(context.Background())
53+
54+
// Set up command
55+
for _, commandHandler := range p.Commands {
56+
handlerMap[commandHandler.Command.Name] = commandHandler.Handler
57+
}
58+
59+
p.Session.AddHandler(func(s *discordgo.Session, i *discordgo.InteractionCreate) {
60+
ctx, cancel := context.WithTimeout(ctx, time.Second*15)
61+
defer cancel()
62+
63+
switch i.Type {
64+
case discordgo.InteractionApplicationCommand:
65+
log.Info().
66+
Str("handlerName", i.ApplicationCommandData().Name).
67+
Msg("invoking command handler")
68+
if h, ok := handlerMap[i.ApplicationCommandData().Name]; ok {
69+
h(ctx, s, i)
70+
}
71+
case discordgo.InteractionMessageComponent:
72+
customID := i.MessageComponentData().CustomID
73+
log.Info().
74+
Str("customID", customID).
75+
Msg("invoking message component handler")
76+
if h, ok := handlerMap[customID]; ok {
77+
h(ctx, s, i)
78+
}
79+
}
80+
})
81+
82+
p.Lifecycle.Append(fx.Hook{OnStart: func(ctx context.Context) error {
83+
for _, v := range p.Commands {
84+
if v.MessageComponent {
85+
continue
86+
}
87+
88+
ccmd, err := p.Session.ApplicationCommandCreate(p.Config.ApplicationID, v.GuildID, &v.Command)
89+
if err != nil {
90+
log.Error().Err(err).Msg("failed to register command")
91+
return err
92+
}
93+
log.Info().
94+
Str("name", ccmd.Name).
95+
Str("id", ccmd.ID).
96+
Msg("registered command")
97+
98+
registeredCommands = append(registeredCommands, ccmd)
99+
}
100+
return nil
101+
},
102+
//OnStop: func(ctx context.Context) error {
103+
// cancel()
104+
// registeredCommands, err := p.Session.ApplicationCommands(p.Config.ApplicationID, "")
105+
// if err != nil {
106+
// log.Error().
107+
// Err(err).
108+
// Msg("could not fetch registered commands")
109+
// return err
110+
// }
111+
// for _, v := range registeredCommands {
112+
// err := p.Session.ApplicationCommandDelete(p.Config.ApplicationID, v.GuildID, v.ID)
113+
// if err != nil {
114+
// log.Error().Err(err).Msg("failed to delete command")
115+
// return err
116+
// }
117+
// log.Debug().
118+
// Str("handlerName", v.Name).
119+
// Str("id", v.ID).
120+
// Msg("deleted command")
121+
// }
122+
// return nil
123+
//}
124+
})
125+
return nil
126+
}

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package discordfx
2+
3+
import (
4+
"github.com/bwmarrin/discordgo"
5+
"github.com/rs/zerolog/log"
6+
"go.uber.org/config"
7+
"go.uber.org/fx"
8+
)
9+
10+
const ConfigurationKey = "discord"
11+
12+
type ConfigParams struct {
13+
fx.In
14+
Config config.Provider
15+
}
16+
17+
type DiscordConfig struct {
18+
ApplicationID string `yaml:"application_id"`
19+
Token string `yaml:"token"`
20+
Intent discordgo.Intent `yaml:"intent"`
21+
}
22+
23+
func NewConfig(p ConfigParams) (cfg DiscordConfig, err error) {
24+
err = p.Config.Get(ConfigurationKey).Populate(&cfg)
25+
if err != nil {
26+
log.Error().Err(err).Msg("failed loading config")
27+
return
28+
}
29+
return
30+
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package discordfx
2+
3+
import (
4+
"go.uber.org/fx"
5+
)
6+
7+
var Module = fx.Options(fx.Provide(NewConfig, NewDiscordSession))

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

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
package discordfx
2+
3+
import (
4+
"github.com/bwmarrin/discordgo"
5+
"github.com/rs/zerolog/log"
6+
"go.uber.org/fx"
7+
)
8+
9+
type NewSessionParams struct {
10+
fx.In
11+
Config DiscordConfig
12+
}
13+
14+
func NewDiscordSession(p NewSessionParams) (*discordgo.Session, error) {
15+
token := "Bot " + p.Config.Token
16+
17+
s, err := discordgo.New(token)
18+
if err != nil {
19+
log.Error().Err(err).Msg("failed to create discord session")
20+
return nil, err
21+
}
22+
23+
if p.Config.Intent > 0 {
24+
s.Identify.Intents = p.Config.Intent
25+
}
26+
27+
if t, err := ParseToken(token); err != nil {
28+
log.Error().Err(err).Msg("failed to parse token")
29+
} else {
30+
log.Info().Str("url", GenerateOAuthURL(t)).Msg("OAuth URL")
31+
}
32+
33+
return s, nil
34+
}
35+
36+
func LogSessionEvents(s *discordgo.Session) {
37+
s.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) {
38+
log.Info().
39+
Interface("event", m).
40+
Msg("message create")
41+
})
42+
//d.session.AddHandler(func(s *discordgo.Session, r *discordgo.Ready) {
43+
// log.Info().
44+
// Str("username", s.State.User.Username).
45+
// Str("discriminator", s.State.User.Discriminator).
46+
// Msg("connected")
47+
//})
48+
//d.session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageEdit) {
49+
// log.Info().
50+
// Interface("event", m).
51+
// Msg("message edit")
52+
//})
53+
//d.session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageDelete) {
54+
// log.Info().
55+
// Interface("event", m).
56+
// Msg("message delete")
57+
//})
58+
//d.session.AddHandler(func(s *discordgo.Session, m *discordgo.PresenceUpdate) {
59+
// log.Info().
60+
// Interface("event", m).
61+
// Msg("presence update")
62+
//})
63+
}

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

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
package discordfx
2+
3+
import (
4+
"errors"
5+
"net/url"
6+
"strconv"
7+
"strings"
8+
)
9+
10+
// Guild represents a Guild ID or Guild URL in the config.
11+
// TODO Make this work properly. Currently requires the user to truncate the channel if they want URL notation.
12+
type Guild string
13+
14+
func (g *Guild) UnmarshalYAML(unmarshal func(interface{}) error) error {
15+
var s string
16+
if err := unmarshal(&s); err != nil {
17+
return err
18+
}
19+
20+
// If it's numeric, just return it.
21+
if _, err := strconv.Atoi(s); err == nil {
22+
*g = Guild(s)
23+
return nil
24+
}
25+
26+
u, err := url.Parse(s)
27+
if err != nil {
28+
return err
29+
}
30+
31+
urlParts := strings.Split(u.Path, "/")
32+
33+
if len(urlParts) < 3 {
34+
return errors.New("invalid guild url")
35+
}
36+
37+
*g = Guild(urlParts[2])
38+
39+
// The result must be numeric.
40+
if _, err := strconv.Atoi(s); err != nil {
41+
return err
42+
}
43+
44+
return nil
45+
}
46+
47+
// Channel represents a Channel ID or Channel URL in the config.
48+
type Channel string
49+
50+
func (c *Channel) UnmarshalYAML(unmarshal func(interface{}) error) error {
51+
var s string
52+
if err := unmarshal(&s); err != nil {
53+
return err
54+
}
55+
56+
// If it's numeric, just return it.
57+
if _, err := strconv.Atoi(s); err == nil {
58+
*c = Channel(s)
59+
return nil
60+
}
61+
62+
u, err := url.Parse(s)
63+
if err != nil {
64+
return err
65+
}
66+
67+
urlParts := strings.Split(u.Path, "/")
68+
69+
if len(urlParts) < 4 {
70+
return errors.New("invalid channel url")
71+
}
72+
73+
*c = Channel(urlParts[3])
74+
75+
// The result must be numeric.
76+
if _, err := strconv.Atoi(string(*c)); err != nil {
77+
return err
78+
}
79+
80+
return nil
81+
}

0 commit comments

Comments
 (0)