hutbot.go

package main

// At a high level, the bot has two concepts: messagers, which produce messages
// and optionally consume responses, and responders, which consume messages and
// optionally produce responses -- consider an IRC messager, that sends a
// message whenever someone speaks in a channel and speaks itself whenever it
// receives a response, and a shell script responder, that runs a shell script
// on each message and sends its output back as a response.
//
// Each messager and responder run in goroutines and communicate on channels;
// every message is dispatched to all responder goroutines and every response
// is dispatched to all messager. So, for example, someone addresses the bot in
// IRC, the IRC messager sends a message that causes a script responder to run
// a script and respond with its output, and the IRC messager then "replies"
// with the script's output in IRC.

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"github.com/thoj/go-ircevent"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

// Command line options.
var botName = flag.String("name", "hutbot", "the bot's nick")
var botPassword = flag.String("password", "", "the bot's password")
var botTLS = flag.Bool("tls", true, "whether to use TLS")
var numWorkers = flag.Int("workers", 4, "process pool size for running commands")

// Messages are created by events (e.g. someone speaking in IRC) and dispatched
// to responders.
type Message struct {
	Messager Messager
	Sender   string
	Channel  string
	Contents string
	Created  time.Time
}

func (message Message) Empty() bool {
	return len(strings.TrimSpace(message.Contents)) <= 0
}

// Responses can be generated by responders (e.g. by running a shell script)
// when they receive messages.
type Response struct {
	Responder Responder
	Message   *Message
	Contents  string
	Target    string // nickname for target, if empty, defaults to channel
	Created   time.Time
}

func (response Response) Empty() bool {
	return len(strings.TrimSpace(response.Contents)) <= 0
}

// Messagers produce messages and consume responses.
type Messager interface {
	Process(chan<- Message, <-chan Response)
}

// StreamMessager produces messages from lines of text on an input stream.
type StreamMessager struct {
	Reader io.Reader
}

func (s *StreamMessager) Process(messages chan<- Message, responses <-chan Response) {
	lines := make(chan string)
	go func() {
		scanner := bufio.NewScanner(s.Reader)
		for scanner.Scan() {
			lines <- scanner.Text()
		}
	}()

	for {
		select {
		case line := <-lines:
			messages <- Message{s, "stdin", "stdin", line, time.Now()}
		case r := <-responses:
			log.Println("[stream]", r)
		}
	}
}

// IRCMessager produces messages from IRC chats.
type IRCMessager struct {
	Server       string
	UseTLS       bool
	Channel      string
	Nick         string
	IdentifyPass string
}

func (i *IRCMessager) callback(cb func(*irc.Event)) func(*irc.Event) {
	return func(event *irc.Event) {
		log.Println("[irc]", event.Code, event, event.Message(), event.Arguments, event.Nick)
		if len(event.Arguments) == 0 || event.Arguments[0] == i.Channel || event.Code == "QUIT" {
			cb(event)
		}
	}
}

func (i *IRCMessager) Connect(conn *irc.Connection) {
	log.Println("[irc] connecting to", i.Server, i.Channel)
	conn.Connect(i.Server)

	if len(i.IdentifyPass) > 0 {
		conn.Privmsg("nickserv", fmt.Sprintf("identify %s %s", i.Nick, i.IdentifyPass))
	}

	log.Println("[irc] joining")
	conn.Join(i.Channel)
}

func (i *IRCMessager) Process(messages chan<- Message, responses <-chan Response) {
	log.Println("[irc] setting up")
	conn := irc.IRC(i.Nick, i.Nick)

	conn.AddCallback("PRIVMSG", i.callback(func(event *irc.Event) {
		messages <- Message{i, event.Nick, event.Arguments[0], event.Message(), time.Now()}
	}))
	conn.AddCallback("JOIN", i.callback(func(event *irc.Event) {
		contents := fmt.Sprintf("%s: irc-join", i.Nick)
		messages <- Message{i, event.Nick, event.Arguments[0], contents, time.Now()}
	}))
	conn.AddCallback("PART", i.callback(func(event *irc.Event) {
		contents := fmt.Sprintf("%s: irc-part", i.Nick)
		messages <- Message{i, event.Nick, event.Arguments[0], contents, time.Now()}
	}))
	conn.AddCallback("QUIT", i.callback(func(event *irc.Event) {
		if event.Nick == *botName {
			time.AfterFunc(10*time.Second, func() {
				log.Println("[irc] re-joining")
				i.Connect(conn)
			})
		} else {
			contents := fmt.Sprintf("%s: irc-quit %s", i.Nick, event.Message())
			messages <- Message{i, event.Nick, i.Channel, contents, time.Now()}
		}
	}))

	conn.UseTLS = i.UseTLS
	i.Connect(conn)

	limiter := NewRateLimiter(1, 5)
	go limiter.Run()

	for response := range responses {
		if response.Empty() {
			continue
		}

		target := response.Target
		if target == "" {
			target = i.Channel
		}
		for _, line := range strings.Split(response.Contents, "\n") {
			limiter.Send(1)
			conn.Privmsg(target, line)
		}
	}
}

// Rate-limits message sending by allocating capacity that drains when messages
// are sent and refills slowly over time, blocking until it's sufficiently
// refilled when empty.
//
// It can be used with any type of message cost -- messages per second, bytes
// per second, etc. -- depending on the values passed to `.Send`.
type RateLimiter struct {
	RefillRate int
	MaxFill    int
	sendchan   chan int
	waitchan   chan int
	level      int
}

// Creates a rate limiter, where `refillRate` is the amount of fill per second
// (in bytes, messages, etc.), and `maxFill` is the amount at which the limiter
// stops filling.
func NewRateLimiter(refillrate int, maxfill int) *RateLimiter {
	return &RateLimiter{
		RefillRate: refillrate,
		MaxFill:    maxfill,
		sendchan:   make(chan int),
		waitchan:   make(chan int),
		level:      maxfill,
	}
}

// Requests to send a message of a certain size or cost, and blocks until we're
// allowed.
func (r *RateLimiter) Send(amount int) int {
	r.sendchan <- amount
	return <-r.waitchan
}

func (r *RateLimiter) refill() {
	r.level += r.RefillRate
	if r.level > r.MaxFill {
		r.level = r.MaxFill
	}
}

// Processes requests to send messages.
func (r *RateLimiter) Run() {
	ticks := time.Tick(time.Second)
	for {
		if r.level <= 0 {
			// If we're "empty", don't take requests (thereby blocking the
			// requester) until we've refilled.
			select {
			case <-ticks:
				r.refill()
			}
		} else {
			// If we're not empty, refill if there's nothing to do, or drain
			// the requested amount.
			select {
			case <-ticks:
				r.refill()
			case amount := <-r.sendchan:
				r.level -= amount
				r.waitchan <- r.level
			}
		}
	}
}

// Invokes a .log script for each non-empty response it gets, and ignores the
// results of the script.
type ResponseLogger struct {
	Pool Pool
}

func (r *ResponseLogger) Process(messages chan<- Message, responses <-chan Response) {
	wd, _ := os.Getwd()
	for response := range responses {
		// Skip empty responses.
		if response.Empty() {
			continue
		}

		// Run .log scripts for each line in the response.
		for _, line := range strings.Split(response.Contents, "\n") {
			env := []string{
				"HUTBOT_EVENT=log_response",
				fmt.Sprintf("HUTBOT_TARGET=%s", response.Target),
				fmt.Sprintf("HUTBOT_CREATED=%d", response.Created.Unix()),
				fmt.Sprintf("HUTBOT_BOT=%s", *botName),
				fmt.Sprintf("HUTBOT_DIR=%s", wd),
				fmt.Sprintf("HUTBOT_RESPONSE=%s", line),
			}
			if response.Message != nil {
				env = append(env,
					fmt.Sprintf("HUTBOT_SENDER=%s", response.Message.Sender),
					fmt.Sprintf("HUTBOT_CHANNEL=%s", response.Message.Channel),
					fmt.Sprintf("HUTBOT_MESSAGE=%s", response.Message.Contents))
			}
			for _, path := range paths(".log") {
				r.Pool.Run(path, "", env, nil)
			}
		}
	}
}

// Invokes a .log script for each non-empty message it gets, and ignores the
// results of the script.
type MessageLogger struct{}

func (m *MessageLogger) Process(pool Pool, messages <-chan Message, responses chan<- Response) {
	wd, _ := os.Getwd()
	for message := range messages {
		if message.Empty() {
			continue
		}
		env := []string{
			"HUTBOT_EVENT=log_message",
			fmt.Sprintf("HUTBOT_SENDER=%s", message.Sender),
			fmt.Sprintf("HUTBOT_CHANNEL=%s", message.Channel),
			fmt.Sprintf("HUTBOT_CREATED=%d", message.Created.Unix()),
			fmt.Sprintf("HUTBOT_BOT=%s", *botName),
			fmt.Sprintf("HUTBOT_DIR=%s", wd),
			fmt.Sprintf("HUTBOT_MESSAGE=%s", message.Contents),
		}
		for _, path := range paths(".log") {
			pool.Run(path, "", env, nil)
		}
	}
}

// Responders consume messages and produce responses.
type Responder interface {
	Process(Pool, <-chan Message, chan<- Response)
}

// Commands are used internally to send data to a pool of workers, so we can
// handle running multiple scripts at once.
type command struct {
	Path     string
	Stdin    string
	Env      []string
	Callback func([]byte, []byte, error)
}

type Pool chan<- command

func (p Pool) Run(path string, stdin string, env []string, callback func([]byte, []byte, error)) {
	p <- command{path, stdin, env, callback}
}

func StartWorkers(num int) Pool {
	commands := make(chan command, 0)
	worker := func(commands <-chan command) {
		for command := range commands {
			stdout, stderr, err := execute(command.Path, command.Stdin, command.Env)
			if command.Callback != nil {
				command.Callback(stdout, stderr, err)
			}
		}
	}
	for i := 0; i < num; i++ {
		go worker(commands)
	}
	return Pool(commands)
}

// PeriodicScript produces unsolicited responses by periodically running
// scripts.
type PeriodicScript struct{}

func (p *PeriodicScript) Process(pool Pool, messages <-chan Message, responses chan<- Response) {
	ticks1Min := time.Tick(time.Minute)
	ticks1Hour := time.Tick(time.Hour)
	ticks1Day := time.Tick(24 * time.Hour)

	wd, _ := os.Getwd()
	env := []string{
		fmt.Sprintf("HUTBOT_BOT=%s", *botName),
		fmt.Sprintf("HUTBOT_DIR=%s", wd),
		"HUTBOT_EVENT=periodic",
	}

	runScripts := func(dir string) {
		for _, path := range paths(dir) {
			executable := path
			pool.Run(path, "", env, func(stdout []byte, stderr []byte, err error) {
				var contents string
				if err == nil {
					contents = strings.TrimRight(string(stdout), " \t\r\n")
				} else {
					contents = fmt.Sprintf("error: %s %s: %s", executable, err, string(stderr))
				}
				responses <- Response{p, nil, contents, "", time.Now()}
			})
		}
	}

	for {
		select {
		case <-messages:
			// Ignore.
		case <-ticks1Min:
			runScripts(".minute")
		case <-ticks1Hour:
			runScripts(".hour")
		case <-ticks1Day:
			runScripts(".day")
		}
	}
}

// CommandScript produces responses whenever messages are addressed to the bot,
// by running scripts whose paths are determined by the content of the message.
type CommandScript struct{}

// Return a slice of the paths of scripts to run based on `name`.
//
// * If `name` is an executable script, include it.
// * If `name` is a directory, include any executable scripts from its
//   immediate children.
func paths(name string) []string {
	// If there's no name, we're done.
	if len(strings.TrimSpace(name)) == 0 {
		return []string{}
	}

	path := "./" + name
	info, err := os.Stat(path)

	// If the path doesn't exist, punt. Otherwise, if it's a directory, inspect
	// its contents. Otherwise, if it's executable, return it.
	if err != nil {
		return []string{}
	} else if info.IsDir() {
		entries, err := ioutil.ReadDir(path)

		// If we can't read the dir, abort.
		if err != nil {
			return []string{}
		}

		// Otherwise, look for executables in the dir.
		result := make([]string, 0)
		for _, entry := range entries {
			if !entry.IsDir() && isExec(entry) {
				result = append(result, path+"/"+entry.Name())
			}
		}
		return result
	} else if isExec(info) {
		return []string{path}
	}
	return []string{}
}

// Is the given file executable by *someone*?
func isExec(f os.FileInfo) bool {
	return (f.Mode()&0111 != 0)
}

// Run the script at `path`, passing it `stdin` and using environment vars
// `env`. Returns stdout and any error that occurred.
func execute(path string, stdin string, env []string) ([]byte, []byte, error) {
	cmd := exec.Command(path)
	cmd.Env = env
	cmd.Stdin = bytes.NewReader([]byte(stdin))
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	return stdout.Bytes(), stderr.Bytes(), err
}

type PathAndTarget struct {
	Path   string
	Target string
}

func (c *CommandScript) Process(pool Pool, messages <-chan Message, responses chan<- Response) {
	pattern := regexp.MustCompile(
		fmt.Sprintf(`%s:\s*([^. \t\r\n]\S*)(\s(.+))?`, *botName))

	for message := range messages {
		var command, args string
		commandFound := false

		match := pattern.FindStringSubmatch(message.Contents)
		if match != nil && len(match) == 4 {
			commandFound = true
			command = match[1]
			args = match[3]
		}
		wd, _ := os.Getwd()
		env := []string{
			"HUTBOT_EVENT=command",
			fmt.Sprintf("HUTBOT_SENDER=%s", message.Sender),
			fmt.Sprintf("HUTBOT_CHANNEL=%s", message.Channel),
			fmt.Sprintf("HUTBOT_CREATED=%d", message.Created.Unix()),
			fmt.Sprintf("HUTBOT_BOT=%s", *botName),
			fmt.Sprintf("HUTBOT_DIR=%s", wd),
			fmt.Sprintf("HUTBOT_COMMAND=%s", command),
			fmt.Sprintf("HUTBOT_ARGS=%s", args),
			fmt.Sprintf("HUTBOT_MESSAGE=%s", message.Contents),
		}

		var pts []PathAndTarget
		if commandFound {
			pts = appendPaths(pts, paths(command), "")
			pts = appendPaths(pts, paths("private/"+command), message.Sender)
			if len(pts) == 0 {
				pts = appendPaths(pts, paths(".missing"), "")
			}
		}
		pts = appendPaths(pts, paths(".all"), "")

		for _, pt := range pts {
			target := pt.Target
			executable := pt.Path
			pool.Run(pt.Path, args, env, func(stdout []byte, stderr []byte, err error) {
				if err == nil {
					contents := strings.TrimRight(string(stdout), " \t\r\n")
					responses <- Response{c, &message, contents, target, time.Now()}
				} else {
					contents := fmt.Sprintf("error: %s %s: %s", executable, err, string(stderr))
					responses <- Response{c, &message, contents, "", time.Now()}
				}
			})
		}
	}
}

func appendPaths(pts []PathAndTarget, paths []string, target string) []PathAndTarget {
	result := pts[:]
	for _, path := range paths {
		result = append(result, PathAndTarget{path, target})
	}
	return result
}

// StartMessager runs the messager in a goroutine and allocates a response
// channel for dispatching reponses to it.
func StartMessager(m Messager, messageChan chan<- Message) chan<- Response {
	responseChan := make(chan Response)
	go m.Process(messageChan, responseChan)
	return responseChan
}

// StartResponder runs the responder in a goroutine and allocates a message
// channel for dispatching messages to it.
func StartResponder(pool Pool, r Responder, responseChan chan<- Response) chan<- Message {
	messageChan := make(chan Message)
	go r.Process(pool, messageChan, responseChan)
	return messageChan
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lmicroseconds)
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] server:port channel\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()
	if len(flag.Args()) != 2 {
		flag.Usage()
		return
	}

	botServer := flag.Arg(0)
	botChannel := flag.Arg(1)

	// Start a worker pool.
	pool := StartWorkers(*numWorkers)

	// Set up messagers.
	messages := make(chan Message, 4096)
	messagers := []Messager{
		&IRCMessager{
			Server:       botServer,
			UseTLS:       *botTLS,
			Nick:         *botName,
			Channel:      botChannel,
			IdentifyPass: *botPassword,
		},
		&StreamMessager{Reader: os.Stdin},
		&ResponseLogger{pool},
	}
	responseChans := []chan<- Response{}
	for _, messager := range messagers {
		responseChans = append(responseChans, StartMessager(messager, messages))
	}

	// Set up responders.
	responses := make(chan Response, 64)
	responders := []Responder{
		&PeriodicScript{},
		&CommandScript{},
		&MessageLogger{},
	}
	messageChans := []chan<- Message{}
	for _, responder := range responders {
		messageChans = append(messageChans, StartResponder(pool, responder, responses))
	}

	// Dispatch responses back to all messagers.
	go func() {
		for response := range responses {
			log.Println("[response]", response)
			for _, messager := range responseChans {
				messager <- response
			}
		}
	}()

	// Dispatch messages to all responders.
	for message := range messages {
		if message.Contents == "__hutbot: quit" {
			break
		}
		log.Println("[message]", message)
		for _, responder := range messageChans {
			responder <- message
		}
	}
}