autogrant_autogrant.go

package autogrant

import (
	"context"
	"fmt"
	"log"
	"strings"

	chime "github.com/nicolai86/aws-chime-sdk"
)

type Poster interface {
	Post(string) (*chime.Message, error)
}

type Command struct {
	Channels []string
}

func (c *Command) Applies(msg, author string) bool {
	return strings.Contains(msg, "requested access to this room")
}

func extractLogin(msg string) (string, string) {
	start := 0
	at := 0
	login := ""
	domain := ""
	for i, c := range msg {
		if c == '(' {
			start = i
			continue
		}

		if c == '@' {
			login = msg[start+1 : i]
			at = i
		}

		if c == ')' {
			domain = msg[at+1 : i]
			break
		}
	}
	return login, domain
}

func (c *Command) Execute(client *chime.Client, p Poster, evt chime.WebsocketPushEvent, msg string) {
	parts := strings.Split(evt.Data.ID, "|")
	roomID := parts[0]
	out, err := client.GetRoom(context.Background(), &chime.GetRoomInput{chime.String(roomID)})
	if err != nil {
		log.Printf("room fetch error: %v", err)
		return
	}
	room := out.Room
	match := false
	for _, name := range c.Channels {
		match = match || strings.Contains(strings.ToLower(*room.Name), strings.ToLower(name))
	}
	if !match {
		log.Printf("Ignoring grant b/c it does not have autogrant enabled: %q", room.Name)
		return
	}
	login, domain := extractLogin(msg)
	if !strings.Contains(domain, "amazon.com") {
		log.Printf("AutoComplete prevented non-amazon login: %v", err)
		return
	}
	// TODO also search by name
	cs, err := client.AutoCompleteContacts(context.Background(), &chime.AutoCompleteContactsInput{
		Query: chime.String(login),
	})
	if err != nil {
		log.Printf("AutoComplete failed: %v", err)
		return
	}
	var cc *chime.Contact
	expectedEmail := fmt.Sprintf("%[email protected]", login)
	for _, c := range *cs.Contacts {
		if *c.Email == expectedEmail {
			cc = &c
			break
		}
	}
	if cc == nil {
		expectedPrefix := fmt.Sprintf("%s@", login)
		for _, c := range *cs.Contacts {
			if strings.HasPrefix(*c.Email, expectedPrefix) {
				cc = &c
				break
			}
		}
	}
	if cc == nil {
		log.Printf("AutoComplete returned %d contacts (expected: 1 match)", len(*cs.Contacts))
		return
	}
	log.Printf("Adding %q to %q", cc.ID, roomID)
	client.AddRoomMember(context.Background(), &chime.AddRoomMemberInput{
		RoomID:    chime.String(roomID),
		ProfileID: cc.ID,
	})
}

main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	vaultapi "github.com/hashicorp/vault/api"
	chime "github.com/nicolai86/aws-chime-sdk"

  "autogrant"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	processedChimeEventCount = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "chime_event_count_processed",
		Help: "Number of chime events processed",
	})
	receivedChimeEventCount = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "chime_event_count_received",
		Help: "Total number of chime events received",
	})
)

func init() {
	prometheus.MustRegister(processedChimeEventCount)
	prometheus.MustRegister(receivedChimeEventCount)
}

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

func NewID() string {
	b := make([]byte, 36) // Same length as uuid v4
	for i := range b {
		b[i] = charset[rand.Intn(len(charset))]
	}
	return string(b)
}

func main() {
	var username, password string
	token := flag.String("vault-token", "", "vault token")
	vaultHost := flag.String("vault-host", "", "vault host")
	config := flag.String("config", "", "chime config to resume sessions")

	nsqHost := flag.String("nsq-host", "", "nsq host to connect to")
	nsqClientCert := flag.String("nsq-client-cert", "", "nsq client cert")
	nsqClientCertKey := flag.String("nsq-client-cert-key", "", "nsq client cert key")
	caCertPath := flag.String("ca-cert", "", "ca cert path")
	_ = nsqHost
	_ = nsqClientCert
	_ = nsqClientCertKey
	_ = caCertPath

	listen := flag.String("listen", ":8899", "interface/ port to listen on")
	flag.Parse()

	http.Handle("/metrics", promhttp.Handler())
	go func() {
		log.Fatal(http.ListenAndServe(*listen, nil))
	}()

	{
		vaultConfig := vaultapi.DefaultConfig()
		vaultConfig.Address = *vaultHost
		client, err := vaultapi.NewClient(vaultConfig)
		if err != nil {
			panic(err)
		}

		client.SetToken(*token)
		c := client.Logical()

		{
			secret, err := c.Read("secret/amazon")
			if err != nil {
				panic(err)
			}
			username = fmt.Sprintf("%[email protected]", secret.Data["login"].(string))
			password = secret.Data["password"].(string)
		}
	}

	if (username == "" || password == "") && config == nil {
		flag.PrintDefaults()
		os.Exit(1)
	}

	opts := []chime.ClientOption{
		chime.WithIDGenerator(NewID),
		chime.WithLogger(os.Stderr),
	}

	client, err := chime.New(opts...)
	if err != nil {
		log.Fatalf(err.Error())
	}

	if err := client.StartSession(nil, username, password); err != nil {
		os.Remove(*config)
		log.Fatalf("failed to start session: %v", err)
	}

	log.Printf("client running")

	ctx := context.Background()
	channels := []string{}
	go func() {

		if err := client.PeriodicRefresh(); err != nil {
			os.Remove(*config)
			log.Fatalf("failed to refresh token %s", err.Error())
			os.Exit(1)
		}
	}()

	rooms, err := client.ListRooms(context.Background(), nil)
	if err != nil {
		log.Fatal(err.Error())
	}

	cmds := []plugins.Command{
		&autogrant.Command{[]string{"golang", "rust", "cdk builders"}},
	}

	for _, room := range rooms.Rooms {
		channels = append(channels, *room.Channel)
	}

	subscriptions, err := client.Subscribe(ctx, &chime.SubscribeInput{
		Subscriptions: channels,
	})
	if err != nil {
		log.Fatalf("websocket connection failed %s", err.Error())
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGTERM)
	signal.Notify(c, syscall.SIGINT)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c
		log.Printf("[chime] shutdown requested")
		client.Close()
		for _, cmd := range cmds {
			v, ok := cmd.(plugins.Stoppable)
			if !ok {
				continue
			}
			<-v.Stop()
		}
		os.Exit(0)
	}()

	go func() {
		<-subscriptions.Context.Done()
		log.Printf("Shutdown complete")
		os.Exit(1)
	}()

	for msg := range subscriptions.Events {
		receivedChimeEventCount.Add(1.0)

		switch msg.Data.Klass {
		case "ConversationMessage":
			{
				processedChimeEventCount.Add(1.0)
				evt := msg.Data.ParsedRecord.(*chime.WebsocketConversationMessageEvent)
				for _, cmd := range cmds {
					if cmd.Applies(evt.Content, evt.Sender) {
						cmd.Execute(client, nil, *msg, evt.Content)
					}
				}
			}
		case "RoomMessage":
			{
				processedChimeEventCount.Add(1.0)
				evt := msg.Data.ParsedRecord.(*chime.WebsocketRoomMessageEvent)

				for _, cmd := range cmds {
					if cmd.Applies(evt.Content, evt.Sender) {
						cmd.Execute(client, nil, *msg, evt.Content)
					}
				}
			}
		default:
			{
			}
		}

	}

	if err := client.Close(); err != nil {
		log.Fatalf(err.Error())
	}
}