Skip to content

Instantly share code, notes, and snippets.

@nicolai86

nicolai86/autogrant_autogrant.go

Created August 27, 2019 16:05
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save nicolai86/ceb63fbedb6ce116a1bd95d024ce2149 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save nicolai86/ceb63fbedb6ce116a1bd95d024ce2149 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. nicolai86 created this gist Aug 27, 2019.
    104 changes: 104 additions & 0 deletions autogrant_autogrant.go
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,104 @@
    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("%s@amazon.com", 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,
    })
    }
    209 changes: 209 additions & 0 deletions main.go
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,209 @@
    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("%s@amazon.com", 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())
    }
    }