A psdoom-aws-ctl - A psdoom-ng compatible command line utility to allow psdoom-ng to kill of EC2 instance in AWS. (see: https://www.youtube.com/watch?v=jJ7AQOBSV1c)

psdoom-aws-ctl.go

package main

import (
	"encoding/gob"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	flag "github.com/spf13/pflag"
)

const (
	cacheFileName  = "instances.gob"
	socketFileName = "instances.sock"
	defaultRegion  = "us-east-1"
)

var (
	profile    string
	region     string
	user       string
	list       bool
	kill       string
	help       bool
	dryRun     bool
	server     bool
	ec2Filters filters
	cacheFile  string
	socketFile string
)

type value struct {
	ID    string
	Name  string
	Type  string
	State int
}

type filters []string

func (t *filters) Set(value string) error {
	*t = append(*t, value)
	return nil
}

func (t *filters) String() string {
	return fmt.Sprint([]string(*t))
}

func (t *filters) Type() string {
	return ""
}

func init() {
	flag.StringVarP(&profile, "profile", "p", "", "Specify name of the AWS profile to use.")
	flag.StringVarP(&region, "region", "r", defaultRegion, "Specify name of the AWS region to use.")
	flag.StringVarP(&user, "user", "u", "", "Specify name of the user to use.")
	flag.BoolVar(&list, "list", false, "List content of the case and exit.")
	flag.StringVar(&kill, "kill", "", "Send given numeric ID to the server to terminate an instance.")
	flag.BoolVarP(&help, "help", "h", false, "Print usage and exist.")
	flag.BoolVar(&dryRun, "dry-run", false, "Test only, do not terminate any instances.")
	flag.BoolVar(&server, "server", false, "Start in the server mode.")
	flag.Var(&ec2Filters, "filters", "A valid one or more filters (see AWS CLI for details).")
	flag.StringVar(&cacheFile, "cache-file", cacheFileName, "Specify cache file to use.")
	flag.StringVar(&socketFile, "socket-file", socketFileName, "Specify Unix socket to use.")
	flag.Parse()
}

func main() {
	instances := make(map[int]value)

	if help {
		flag.Usage()
		os.Exit(0)
	}

	if kill != "" {
		addr, err := net.ResolveUnixAddr("unix", socketFile)
		if err != nil {
			panic(err)
		}

		conn, err := net.DialUnix("unix", nil, addr)
		if err != nil {
			panic(err)
		}
		defer func() {
			if err := conn.Close(); err != nil {
				panic(err)
			}
		}()

		_, err = conn.Write([]byte(strings.TrimSpace(kill)))
		if err != nil {
			panic(err)
		}
		return
	}

	sess := session.Must(session.NewSessionWithOptions(session.Options{
		Profile: profile,
	}))

	region := aws.StringValue(sess.Config.Region)
	if region == "" {
		region = defaultRegion
	}

	svc := ec2.New(sess, aws.NewConfig().WithRegion(region))

	if server {
		fmt.Println("Server started.")

		c := make(chan os.Signal)

		signal.Notify(c, os.Interrupt, syscall.SIGTERM)
		go func() {
			<-c
			err := os.Remove(socketFile)
			if err != nil {
				panic(err)
			}
			os.Exit(1)
		}()

		f, err := os.Open(cacheFile)
		if err != nil {
			panic(err)
		}

		dec := gob.NewDecoder(f)

		err = dec.Decode(&instances)
		if err != nil {
			panic(err)
		}
		if err := f.Close(); err != nil {
			panic(err)
		}

		addr, err := net.ResolveUnixAddr("unix", socketFile)
		if err != nil {
			panic(err)
		}

		l, err := net.ListenUnix("unix", addr)
		if err != nil {
			panic(err)
		}

		var mutex sync.Mutex

		for {
			conn, err := l.AcceptUnix()
			if err != nil {
				fmt.Println(err)
				continue
			}

			go func(c *net.UnixConn) {
				r, err := ioutil.ReadAll(c)
				if err != nil {
					fmt.Println(err)
					return
				}
				defer func() {
					if err := c.Close(); err != nil {
						fmt.Println(err)
						return
					}
				}()

				n, err := strconv.Atoi(strings.TrimSpace(string(r)))
				if err != nil {
					fmt.Println(err)
					return
				}

				t, ok := instances[n]
				if ok {
					prefix := ""
					if dryRun {
						prefix = "(dry run) "
					}
					fmt.Printf("%sTerminating instance: %s (Name=%s Type=%s)\n", prefix, t.ID, t.Name, t.Type)

					if !dryRun {
						_, err := svc.TerminateInstances(&ec2.TerminateInstancesInput{
							InstanceIds: []*string{aws.String(t.ID)},
						})
						if err != nil {
							fmt.Println(err)
							return
						}
					}

					mutex.Lock()
					defer mutex.Unlock()

					delete(instances, n)

					f, err = os.Create(cacheFile)
					if err != nil {
						fmt.Println(err)
						return
					}
					defer func() {
						if err := f.Close(); err != nil {
							fmt.Println(err)
							return
						}
					}()

					enc := gob.NewEncoder(f)

					err = enc.Encode(instances)
					if err != nil {
						fmt.Println(err)
						return
					}

				}
			}(conn)
		}
	}

	if list {
		f, err := os.Open(cacheFile)
		if err != nil {
			panic(err)
		}
		defer func() {
			if err := f.Close(); err != nil {
				panic(err)
			}
		}()

		s, err := f.Stat()
		if err != nil {
			panic(err)
		}
		if s.Size() == 0 {
			os.Exit(0)
		}

		dec := gob.NewDecoder(f)

		err = dec.Decode(&instances)
		if err != nil {
			panic(err)
		}

		if user == "" {
			user = os.Getenv("USER")
		}

		for k, v := range instances {
			fmt.Printf("%s %d %s %d\n", user, k, v.Name, v.State)
		}

		os.Exit(0)
	}

	params := &ec2.DescribeInstancesInput{}

	if len(ec2Filters) > 0 {
		var pairs [][]string
		for _, filter := range ec2Filters {
			if strings.HasPrefix(filter, "Name=") {
				var parts []string
				for _, s := range strings.SplitN(filter, ",", 2) {
					parts = append(parts, strings.Split(s, "=")...)
				}
				pairs = append(pairs, []string{parts[1], parts[3]})
			} else {
				for _, s := range strings.Split(filter, ",") {
					pairs = append(pairs, strings.Split(s, "="))
				}
			}
		}

		var filters []*ec2.Filter
		for _, pair := range pairs {
			if len(pair) < 2 || pair[0] == "" || pair[1] == "" {
				continue
			}
			filters = append(filters, &ec2.Filter{
				Name:   aws.String(pair[0]),
				Values: aws.StringSlice(strings.Split(pair[1], ",")),
			})
		}
		params.Filters = filters
	}

	fmt.Print("Refreshing cache...")

	resp, err := svc.DescribeInstances(params)
	if err != nil {
		panic(err)
	}

	var name string
	var state string

	i := 1
	for _, res := range resp.Reservations {
		for _, inst := range res.Instances {
			state = aws.StringValue(inst.State.Name)
			if state == "terminated" && state != "stopped" && state != "running" {
				continue
			}

			v := value{
				ID:    aws.StringValue(inst.InstanceId),
				Type:  aws.StringValue(inst.InstanceType),
				State: 0,
			}

			if len(inst.Tags) > 0 {
				for _, tag := range inst.Tags {
					if aws.StringValue(tag.Key) == "Name" {
						name = aws.StringValue(tag.Value)
						break
					}
				}
			}

			v.Name = v.ID
			if name != "" {
				v.Name = name
			}

			if state == "running" {
				v.State = 1
			}

			instances[i] = v
			i++
		}
	}

	f, err := os.Create(cacheFile)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			panic(err)
		}
	}()

	enc := gob.NewEncoder(f)

	err = enc.Encode(instances)
	if err != nil {
		panic(err)
	}

	fmt.Println(" OK.")
}

psdoom-ng.sh

#!/bin/bash

set +e

export AWS_REGION='us-east-1'
export AWS_SECRET_KEY=''
export AWS_SECRET_ACCESS_KEY=''

export PSDOOMPSCMD='./psdoom-aws-ctl --list'
export PSDOOMRENICECMD='true'
export PSDOOMKILLCMD='./psdoom-aws-ctl --kill'

rm -f \
	instances.gob \
	instances.sock

./psdoom-aws-ctl
./psdoom-aws-ctl --server &
server=$!

./psdoom-ng -nomouse -window -puser "$USER" -episode 1 -godstart "$@"

{
	kill -0 "$server" && kill -TERM "$server"
} &> /dev/null