package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"syscall"
	"time"
)

var Error *log.Logger
var Warning *log.Logger

func init() {
	log.SetFlags(log.Lmicroseconds | log.Lshortfile)
	logFile, err := os.Create("main.log")
	if err != nil {
		log.Printf("Error creating log file: %v", err)
		os.Exit(1)
	}
	logger := io.MultiWriter(os.Stdout, logFile)
	log.SetOutput(logger)

	Error = log.New(io.MultiWriter(logFile, os.Stderr, os.Stdout),
		fmt.Sprintf("%sERROR:%s ", "\033[0;101m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)
	Warning = log.New(io.MultiWriter(logFile, os.Stdout),
		fmt.Sprintf("%sWarning:%s ", "\033[0;93m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)
}

func main() {
	forbidden, exists := os.LookupEnv("HITMAN_FORBIDDEN")
	if !exists {
		Error.Println("HITMAN_FORBIDDEN environment variable not set")
		log.Printf("Please set to a comma separated list of process names to forbid")
		return
	}
	delay := time.Duration(3) * time.Second
	scanDelay, exists := os.LookupEnv("HITMAN_SCAN_DELAY")
	if !exists {
		log.Printf("No scan delay is set, defaulting to %vs", delay.Seconds())
		log.Printf("Set HITMAN_SCAN_DELAY to change this")
	} else {
		var err error
		delay, err = time.ParseDuration(scanDelay)
		if err != nil {
			Error.Printf("Error parsing scan delay: %v", err)
			return
		}
	}

	procs := strings.Split(forbidden, ",")

	for {
		log.Printf("Running")
		procmap, err := BuildProcessMap()
		if err != nil {
			Error.Printf("Error building process map: %v", err)
			return
		}

		for _, proc := range procs {
			log.Printf("Checking %s", proc)
			res, ok := procmap.findByName(proc)
			if ok {
				log.Printf("Forbidden process %s found (x%d)", proc, len(res))
				for _, node := range res {
					log.Printf("Killing forbidden process %d", node.Proc.ProcessID)
					err := Kill(node.Proc.ProcessID)
					if err != nil {
						Error.Printf("Error terminating process %d: %v", node.Proc.ProcessID, err)
					}
				}
				return
			}
		}
		time.Sleep(delay)
	}
}

func Kill(pid uint32) error {
	handle, err := syscall.OpenProcess(syscall.PROCESS_TERMINATE, false, uint32(pid))
	if err != nil {
		return fmt.Errorf("error opening process: %v", err)
	}
	defer syscall.CloseHandle(handle)

	err = syscall.TerminateProcess(handle, 7172)
	if err != nil {
		return fmt.Errorf("error terminating process: %v", err)
	}

	return nil
}