package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"os"
)

var debug = false

func encrypt(s string) []byte {
	key := byte(171)
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.BigEndian, uint32(len(s)))

	// Iterate through the string and apply the XOR cipher
	for i := 0; i < len(s); i++ {
		a := key ^ s[i]
		key = a
		buf.WriteByte(a)
	}

	return buf.Bytes()
}

func decrypt(ciphertext []byte) string {
	if len(ciphertext) < 4 {
		return ""
	}
	payload := ciphertext[4:]
	key := byte(171)
	result := make([]byte, len(payload))

	for i, b := range payload {
		a := key ^ b
		key = b
		result[i] = a
	}
	return string(result)
}

func sendPayload(ip, command string) (string, error) {
	conn, err := net.Dial("tcp", ip+":9999")
	if err != nil {
		return "", fmt.Errorf("failed to connect: %w", err)
	}
	defer conn.Close()

	encPayload := encrypt(command)
	if _, err := conn.Write(encPayload); err != nil {
		return "", fmt.Errorf("failed to send data: %w", err)
	}

	response, err := io.ReadAll(conn)
	if err != nil {
		return "", fmt.Errorf("failed to get response: %w", err)
	}

	decryptResponse := decrypt(response[4:])
	return decryptResponse, nil
}

func validateIp(ip string) bool {
	return net.ParseIP(ip) != nil
}

func main() {
	nea := "Not enough arguments provided"
	if len(os.Args) > 1 {
		userIp := os.Args[1]
		var cmd string
		if validateIp(userIp) {
			if len(os.Args) > 2 {
				switch os.Args[2] {
				case "on", "1":
					cmd = `{"system":{"set_relay_state":{"state":1}}}`
				case "off", "0":
					cmd = `{"system":{"set_relay_state":{"state":0}}}`
				default:
					fmt.Println("Invalid argument")
					os.Exit(1)
				}
			} else {
				fmt.Println(nea)
				os.Exit(1)
			}
			response, err := sendPayload(userIp, cmd)
			if err != nil {
				log.Fatal(err)
			}
			if debug {
				fmt.Printf("response: %s\n", response)
			}
		} else {
			fmt.Println("Invalid IP provided")
			os.Exit(1)
		}
	} else {
		fmt.Println(nea)
		fmt.Printf("Usage %s {ip} {on|off}\n", os.Args[0])
		os.Exit(1)
	}
	os.Exit(0)
}