msg.kwolek.io/main.go

package main

import (
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/gliderlabs/ssh"
	"github.com/joho/godotenv"

	terminal "golang.org/x/term"
)

type Message struct {
	Content   string
	Username  string
	PK        ssh.PublicKey
	PKShort   string
	Timestamp time.Time
}

type MessageList struct {
	list []Message
	size int
}

func NewMessageList(size int) *MessageList {
	ml := &MessageList{size: size}
	return ml
}

func (ml *MessageList) Add(m Message) {
	if len(ml.list) == ml.size {
		ml.list = ml.list[1:]
	}
	ml.list = append(ml.list, m)
}

func (ml *MessageList) List() []Message {
	return ml.list
}

func (ml *MessageList) Reverse() []Message {
	r := make([]Message, len(ml.list))
	l := len(ml.list) - 1
	for i, m := range ml.list {
		r[l-i] = m
	}
	return r
}

func CheckString(s string) bool {
	for i := 0; i < len(s); i++ {
		if s[i] < 32 || s[i] > 126 {
			return false
		}
	}
	return true
}

func main() {
	godotenv.Load()
	sshKey := ssh.HostKeyFile(os.Getenv("SSH_KEY"))
	sshAddr := os.Getenv("SSH_ADDR")
	httpAddr := os.Getenv("HTTP_ADDR")
	siteName := os.Getenv("SITE_NAME")
	sshRemoteAddr := os.Getenv("SSH_REMOTE_ADDR")

	m := NewMessageList(10)
	pkAuth := ssh.PublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
		return true
	})
	ssh.Handle(func(s ssh.Session) {
		term := terminal.NewTerminal(s, "> ")
		user := s.User()
		term.Write(term.Escape.Red)
		defer term.Write(term.Escape.Reset)
		if !CheckString(user) {
			io.WriteString(term, "Sorry, your username needs to be ASCII (without funky characters).\n")
			return
		}
		pk := s.PublicKey()
		if pk == nil {
			io.WriteString(s, "Hi! Please come back with a public key :)\n")
			return
		}
		term.Write(term.Escape.Reset)

		pkMarshalled := pk.Marshal()
		pkHex := hex.EncodeToString(pkMarshalled)
		pkHexShort := pkHex[len(pkHex)-7:]

		fmt.Fprintf(s,
			"%sHello, %s%s%s (%s)%s\n",
			term.Escape.Blue,
			term.Escape.Cyan, user,
			term.Escape.Blue, pkHexShort,
			term.Escape.Reset)

		fmt.Fprintf(s, "%sRun %s`w`%s to write a message,  %s`p`%s to print existing messages and %s`q`%s to quit.%s\n",
			term.Escape.Blue,
			term.Escape.Cyan, term.Escape.Blue,
			term.Escape.Cyan, term.Escape.Blue,
			term.Escape.Cyan, term.Escape.Blue,
			term.Escape.Reset)

		for {
			term.SetPrompt(fmt.Sprintf("\n%s>%s ", term.Escape.Blue, term.Escape.Reset))
			cmd, err := term.ReadLine()
			if err == io.EOF {
				cmd = "q"
				err = nil
			}
			if err != nil {
				return
			}

			switch strings.TrimSpace(cmd) {
			case "w":
				term.SetPrompt(fmt.Sprintf("%sMessage:%s ", term.Escape.Blue, term.Escape.Reset))
				msg, err := term.ReadLine()
				if err == io.EOF {
					continue
				}
				fmt.Fprintf(term, "%sYour message:%s %s%s\n", term.Escape.Blue, term.Escape.White, msg, term.Escape.Reset)
				if !CheckString(msg) {
					fmt.Fprintf(term, "Sorry, your message contains non-ASCII or control characters.\n")
					continue
				}
				term.SetPrompt(fmt.Sprintf(
					"%sPublish? (%syes%s/%sno%s):%s ",
					term.Escape.Blue,
					term.Escape.Green, term.Escape.Blue,
					term.Escape.Green, term.Escape.Blue,
					term.Escape.Reset,
				))
				cmd, _ := term.ReadLine()
				if strings.ToLower(strings.TrimSpace(cmd)) == "yes" {
					m.Add(Message{
						Content:   msg,
						Username:  user,
						PK:        pk,
						PKShort:   pkHexShort,
						Timestamp: time.Now(),
					})
					fmt.Fprintf(term, "%sPublished!%s\n", term.Escape.Green, term.Escape.Reset)
				} else {
					fmt.Fprintf(term, "%sOkay, not publishing.%s\n", term.Escape.Blue, term.Escape.Reset)
				}
			case "p":
				for _, msg := range m.List() {
					fmt.Fprintf(term,
						"%s%s <%s%s%s (%s)%s> %s%s%s\n",
						term.Escape.Blue,
						msg.Timestamp.Format("2006-01-02 15:04:05 MST"),
						term.Escape.Cyan, msg.Username, term.Escape.Blue,
						msg.PKShort, term.Escape.Blue,
						term.Escape.White, msg.Content, term.Escape.Reset)
				}
			case "q":
				fmt.Fprintf(term, "Bye!\n")
				return
			default:
				fmt.Fprintln(term, "Unknown command.")
			}
		}
	})

	web := Web{
		Messages: m,
		SiteName: siteName,
		SSHAddr: sshRemoteAddr,
	}
	go web.ListenAndServe(httpAddr)
	log.Fatal(ssh.ListenAndServe(sshAddr, nil, pkAuth, sshKey))
}