path: root/main.go

package main

import (
	"container/heap"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/radareorg/r2pipe-go"
	"github.com/sirupsen/logrus"
)

// Config defines the meta config
type Config struct {

	// Arch defines the architecture the battle should run in
	Arch string

	// Bits defines the bitness
	Bits int

	// Memsize defines the arena size
	Memsize int

	// MaxProgSize defines the maximal bot size
	MaxProgSize int

	// Bots defines a list of bots to take part in the battle
	Bots []Bot

	// AmountOfBots defines the amount of bots taking part in the tournament
	AmountOfBots int

	// RandomOffsets defines the offset in memory where the bots should be placed
	RandomOffsets []int

	// GameRoundTime defines the length of a gameround
	GameRoundDuration time.Duration
}

// Bot defines a bot
type Bot struct {

	// Path defines the path to the source of the bot
	Path string

	// Code holds the assembled code of the bot as generated by rasm2
	Code string

	// Addr defines the initial address the bot is placed at
	Addr int

	// Regs defines the state of the registers of the bot
	// It is used to store the registers after each round and restore them in the
	// next round when the bot's turn has come
	Regs string
}

func parseConfig() Config {
	arch := flag.String("arch", "x86", "bot architecture (mips|arm|x86)")
	bits := flag.Int("bits", 32, "bot bitness (8|16|32|64)")
	maxProgSize := flag.Int("maxProgSize", 64, "the maximum bot size")
	memPerBot := flag.Int("memPerBot", 512, "the amount of memory each bot should add to the arena")
	gameRoundDuration := flag.Duration("t", 250*time.Millisecond, "The duration of a round")

	v := flag.Bool("v", false, "info")
	vv := flag.Bool("vv", false, "debug")
	vvv := flag.Bool("vvv", false, "trace")

	flag.Parse()

	if *v == true {
		logrus.SetLevel(logrus.InfoLevel)
	} else if *vv == true {
		logrus.SetLevel(logrus.DebugLevel)
	} else if *vvv == true {
		logrus.SetLevel(logrus.TraceLevel)
	} else {
		logrus.SetLevel(logrus.WarnLevel)
	}

	// parse all trailing command line arguments as path to bot sourcecode
	amountOfBots := flag.NArg()

	memsize := amountOfBots * *memPerBot

	logrus.WithFields(logrus.Fields{
		"mem per bot":  *memPerBot,
		"amountOfBots": amountOfBots,
		"memsize":      memsize,
	}).Infof("Loaded config")

	// define a config to return
	config := Config{
		Arch:              *arch,
		Bits:              *bits,
		Memsize:           memsize,
		MaxProgSize:       *maxProgSize,
		AmountOfBots:      amountOfBots,
		GameRoundDuration: *gameRoundDuration,
	}

	return config
}

// define bots defines the bots given via command line arguments
func defineBots(config *Config) {

	logrus.Info("Defining the bots")

	// define a list of bots by parsing the command line arguments one by one
	var bots []Bot
	for i := 0; i < config.AmountOfBots; i++ {
		bot := Bot{
			Path: flag.Arg(i),
		}
		bots = append(bots, bot)
	}

	config.Bots = bots
}

func r2cmd(r2p *r2pipe.Pipe, input string) string {

	logrus.Tracef("> %s", input)

	// send a command
	buf1, err := r2p.Cmd(input)
	if err != nil {
		panic(err)
	}

	// return the result of the command as a string
	return buf1
}

func buildBots(config *Config) {

	logrus.Info("Building all bots")

	// build all the bots
	for i := 0; i < config.AmountOfBots; i++ {
		buildBot(i, config)
	}
}

// buildBot builds the bot located at the given path.
func buildBot(i int, config *Config) {

	logrus.Debugf("Building bot %d", i)

	// open radare without input for building the bot
	r2p1, err := r2pipe.NewPipe("--")
	if err != nil {
		panic(err)
	}
	defer r2p1.Close()

	// Compile a warrior using rasm2
	//
	// This uses the given architecture, the given bitness and the given path in
	// rasm2 to compile the bot
	botPath := config.Bots[i].Path
	radareCommand := fmt.Sprintf("rasm2 -a %s -b %d -f %s", config.Arch, config.Bits, botPath)
	botCode := r2cmd(r2p1, radareCommand)

	config.Bots[i].Code = botCode
}

// init initializes the arena
func initArena(config *Config) *r2pipe.Pipe {

	logrus.Info("Initializing the arena")
	logrus.Debugf("Allocating %d bytes of memory...", config.Memsize)

	// allocate memory
	r2p, err := r2pipe.NewPipe(fmt.Sprintf("malloc://%d", config.Memsize))
	if err != nil {
		panic(err)
	}

	logrus.Info("Memoy successfully allocated")

	// define the architecture and the bitness
	_ = r2cmd(r2p, fmt.Sprintf("e asm.arch = %s", config.Arch))
	_ = r2cmd(r2p, fmt.Sprintf("e asm.bits = %d", config.Bits))

	// enable colors
	// _ = r2cmd(r2p, "e scr.color = 0")
	_ = r2cmd(r2p, "e scr.color = 3")
	_ = r2cmd(r2p, "e scr.color.args = true")
	_ = r2cmd(r2p, "e scr.color.bytes = true")
	_ = r2cmd(r2p, "e scr.color.grep = true")
	_ = r2cmd(r2p, "e scr.color.ops = true")
	_ = r2cmd(r2p, "e scr.bgfill = true")
	_ = r2cmd(r2p, "e scr.color.pipe = true")
	_ = r2cmd(r2p, "e scr.utf8 = true")

	// hex column width
	_ = r2cmd(r2p, "e hex.cols = 32")

	// initialize ESIL VM state
	logrus.Debug("Initializing the ESIL VM")
	_ = r2cmd(r2p, "aei")

	// initialize ESIL VM stack
	logrus.Debug("Initializing the ESIL Stack")
	_ = r2cmd(r2p, "aeim")

	// return the pipe
	return r2p
}

// a heap of integers (adapted from the container/heap documentation)
type IntHeap []int

func (h IntHeap) Len() int		{ return len(h) }
func (h IntHeap) Less(i, j int) bool	{ return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)		{ h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x any) {
	*h = append(*h, x.(int))
}

func (h *IntHeap) Pop() any {
	n := len(*h)
	x := (*h)[n - 1]
	*h = (*h)[0 : n - 1]
	return x
}

// genRandomOffsets returns random offsets for all bots
// This is used to get the offset the bots are initially placed in
func genRandomOffsets(config *Config) {

	logrus.Info("Generating random bot offsets")

	// seed the random number generator
	rand.Seed(time.Now().UTC().UnixNano())

	nbots := len(config.Bots)

	// calculate final occupied and remaining free space
	needed := 0
	for i := 0; i < nbots; i++ {
		needed += len(config.Bots[i].Code)
	}
	if needed > config.Memsize {
		panic("total bot size exceeds available memory")
	}
	free := config.Memsize - needed

	// pick nbots random places among chunks of free memory.
	// we consider as possible places the spaces between individual free
	// bytes of memory and pick one uniformly at random for each bot.
	// the results are stored in a heap so we can easily go through them
	// from first to last later.
	// it is okay to pick the same place multiple times, which just means
	// that the corresponding bots will be adjacent.
	places := &IntHeap{}
	for i := 0; i < nbots; i++ {
		// there are free + 1 possible places, one in front of each
		// free byte, plus one after the last.
		heap.Push(places, rand.Intn(free + 1))
	}

	// randomly choose the order in which the bots should appear
	perm := rand.Perm(nbots)

	// calculate the resulting bot offsets
	offsets := make([]int, nbots)
	occupied := 0;				// space allocated so far
	for i := 0; i < nbots; i++ {
		bot := perm[i]
		size := len(config.Bots[bot].Code)
		place := heap.Pop(places).(int)

		offsets[bot] = place + occupied
		if offsets[bot] + size > config.Memsize {
			panic("BUG: bot placed out of bounds")
		}

		occupied += size
	}

	// debug print all offsets
	var fields0 logrus.Fields = make(logrus.Fields)
	for i := 0; i < len(offsets); i++ {
		fields0[fmt.Sprintf("%d", i)] = offsets[i]
	}
	logrus.WithFields(fields0).Debug("Offsets")

	config.RandomOffsets = offsets
}

// place the bot in the arena at the given address
func placeBot(r2p *r2pipe.Pipe, bot Bot, address int) {
	_ = r2cmd(r2p, fmt.Sprintf("wx %s @ %d", bot.Code, address))
}

func placeBots(r2p *r2pipe.Pipe, config *Config) {

	logrus.Info("Placing the bots in the arena")

	// place each bot in the arena
	for bot := 0; bot < len(config.Bots); bot++ {

		// get the address where the bot should be placed
		address := config.RandomOffsets[bot]

		// Place the bot in the arena
		logrus.Debugf("[i] Placing bot %d at %d", bot, address)
		placeBot(r2p, config.Bots[bot], address)

		logrus.Debugf("\n%s", r2cmd(r2p, fmt.Sprintf("pd 0x8 @ %d", address)))

		// store the initial address of the bot in the according struct field
		config.Bots[bot].Addr = address

		// define the instruction point and the stack pointer
		_ = r2cmd(r2p, fmt.Sprintf("aer PC=%d", config.Bots[bot].Addr))
		_ = r2cmd(r2p, fmt.Sprintf("aer SP=SP+%d", config.Bots[bot].Addr))

		// dump the registers of the bot for being able to switch inbetween them
		// This is done in order to be able to play one step of each bot at a time,
		// but sort of in parallel
		initialRegisers := strings.Replace(r2cmd(r2p, "aerR"), "\n", ";", -1)
		config.Bots[bot].Regs = initialRegisers
	}
}

func defineErrors(r2p *r2pipe.Pipe) {
	// handle errors in esil
	_ = r2cmd(r2p, "e cmd.esil.todo=f theend=1")
	_ = r2cmd(r2p, "e cmd.esil.trap=f theend=1")
	_ = r2cmd(r2p, "e cmd.esil.intr=f theend=1")
	_ = r2cmd(r2p, "e cmd.esil.ioer=f theend=1")
	_ = r2cmd(r2p, "f theend=0")
}

// StepIn steps in and stores the state of the registers for the given bot
func stepIn(r2p *r2pipe.Pipe) {
	_ = r2cmd(r2p, "aes")
}

// switchPlayer returns the id of the next Player
func switchPlayer(r2p *r2pipe.Pipe, currentPlayer int, config *Config) int {

	// calculate the index of the nextPlayer
	nextPlayer := (currentPlayer + 1) % config.AmountOfBots

	return nextPlayer
}

func arena(r2p *r2pipe.Pipe, config *Config, id, gen int) string {
	var res string = ""

	// clear the screen
	res += "\x1b[2J\x1b[0;0H"
	// res += fmt.Sprintf("%s\n", r2cmd(r2p, "?eg 0 0"))

	// print some general information such as the current user and the round the
	// game is in
	ip := fmt.Sprintf("%s\n", r2cmd(r2p, "aer~eip"))
	res += fmt.Sprintf("Round: %d \t\t User: %d \t\t ip: %s\n", gen, id, ip)

	// print the memory space
	res += fmt.Sprintf("%s\n", r2cmd(r2p, "pxa 0x400 @ 0"))
	// res += fmt.Sprintf("%s\n", r2cmd(r2p, fmt.Sprintf("pd 0x10 @ %d", config.Bots[id].Addr)))

	// res += fmt.Sprintf("%s\n", r2cmd(r2p, "prc 0x200 @ 0"))

	return res
}

// runGame actually runs the game (surprise!)
func runGame(r2p *r2pipe.Pipe, config *Config) {

	// start the competition
	var botid int = 0
	var round int = 0
	for true {

		// load the registers
		r2cmd(r2p, config.Bots[botid].Regs)

		// Step
		stepIn(r2p)

		// store the regisers
		registers := r2cmd(r2p, "aerR")
		registersStripped := strings.Replace(registers, "\n", ";", -1)
		config.Bots[botid].Regs = registersStripped

		logrus.Info(arena(r2p, config, botid, round))

		if dead(r2p, botid) == true {
			logrus.Warnf("DEAD (round %d)", round)
			os.Exit(1)
		}

		// switch players, if the new botid is 0, a new round has begun
		botid = switchPlayer(r2p, botid, config)
		if botid == 0 {
			round++
		}

		// sleep only a partial of the total round time, as a round is made up of
		// the movements of multiple bots
		time.Sleep(config.GameRoundDuration / time.Duration(config.AmountOfBots))
	}
}

func dead(r2p *r2pipe.Pipe, botid int) bool {
	status := r2cmd(r2p, "?v 1+theend")

	if status != "" && status != "0x1" {
		logrus.Warnf("[!] Bot %d has died", botid)
		return true
	}
	return false
}

func main() {
	fmt.Println("hi")

	config := parseConfig()
	defineBots(&config)
	buildBots(&config)
	genRandomOffsets(&config)

	// initialize the arena (allocate memory + initialize the ESIL VM & stack)
	r2p := initArena(&config)

	// place the bots in the arena
	placeBots(r2p, &config)

	// if an error occurs (interrupt, ioerror, trap, ...), the ESIL VM should set
	// a flag that can be used to determine if a player has died
	defineErrors(r2p)

	// run the actual game
	runGame(r2p, &config)

	r2p.Close()
}