From 2aff4d937adfb3c487e33468777b10b92af17bbc Mon Sep 17 00:00:00 2001
From: Dreaded_X <tim@huizinga.dev>
Date: Fri, 18 Nov 2022 04:05:13 +0100
Subject: [PATCH] Made the kasa integration more robust and added GetState
 function to get the current state

---
 integration/kasa/kasa.go    | 82 ++++++++++++++++++++++++++++---------
 integration/kasa/message.go | 26 ++++++++++++
 integration/kasa/reply.go   | 19 ---------
 3 files changed, 89 insertions(+), 38 deletions(-)
 create mode 100644 integration/kasa/message.go
 delete mode 100644 integration/kasa/reply.go

diff --git a/integration/kasa/kasa.go b/integration/kasa/kasa.go
index 6b9309c..f339782 100644
--- a/integration/kasa/kasa.go
+++ b/integration/kasa/kasa.go
@@ -5,9 +5,13 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"log"
 	"net"
 )
 
+// This implementation is based on:
+// https://www.softscheck.com/en/blog/tp-link-reverse-engineering/
+
 func encrypt(data []byte) []byte {
 	var key byte = 171
 	buf := new(bytes.Buffer)
@@ -22,18 +26,23 @@ func encrypt(data []byte) []byte {
 	return buf.Bytes()
 }
 
-func decrypt(data []byte) string {
+func decrypt(data []byte) ([]byte, error) {
 	var key byte = 171
-	buf := new(bytes.Buffer)
-	binary.Write(buf, binary.BigEndian, uint32(len(data)))
 
-	for _, c := range data {
-		a := key ^ c
-		key = c
-		buf.WriteByte(a)
+	if len(data) < 4 {
+		return nil, fmt.Errorf("Array has a minumun size of 4")
 	}
 
-	return string(buf.Bytes())
+	size := binary.BigEndian.Uint32(data[0:4])
+	buf := make([]byte, size)
+
+	for i := 0; i < int(size); i++ {
+		a := key ^ data[i+4]
+		key = data[i+4]
+		buf[i] = a
+	}
+
+	return buf, nil
 }
 
 
@@ -45,44 +54,79 @@ func New(ip string) Kasa {
 	return Kasa{ip}
 }
 
-func (kasa *Kasa) sendCmd(cmd cmd) {
+func (kasa *Kasa) sendCmd(cmd cmd) (reply, error) {
 	con, err := net.Dial("tcp", fmt.Sprintf("%s:9999", kasa.ip))
 	if err != nil {
-		panic(err)
+		return reply{}, err
 	}
 
 	defer con.Close()
 
 	b, err := json.Marshal(cmd)
 	if err != nil {
-		panic(err)
+		return reply{}, err
 	}
 
 	_, err = con.Write(encrypt(b))
 	if err != nil {
-		panic(err)
+		return reply{}, err
 	}
 
 	resp := make([]byte, 2048)
 	_, err = con.Read(resp)
 	if err != nil {
-		panic(err)
+		return reply{}, err
+	}
+
+	d, err := decrypt(resp)
+	if err != nil {
+		return reply{}, err
 	}
 
 	var reply reply
-	json.Unmarshal(resp, &reply)
-
-	if reply.System.SetRelayState.ErrCode != 0 {
-		fmt.Println(reply)
-		fmt.Println(resp)
+	err = json.Unmarshal(d, &reply)
+	if err != nil {
+		return reply, err
 	}
+
+	return reply, err
 }
 
 func (kasa *Kasa) SetState(on bool) {
 	var cmd cmd
+	cmd.System.SetRelayState = &SetRelayState{State: 0}
 	if on {
 		cmd.System.SetRelayState.State = 1
 	}
 
-	kasa.sendCmd(cmd)
+	reply, err := kasa.sendCmd(cmd)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+
+	if reply.System.SetRelayState.ErrCode != 0 {
+		log.Printf("Failed to set relay state, error: %d\n", reply.System.SetRelayState.ErrCode)
+	}
 }
+
+func (kasa *Kasa) GetState() bool {
+	cmd := cmd{}
+
+	cmd.System.GetSysinfo = &GetSysinfo{}
+	
+	reply, err := kasa.sendCmd(cmd)
+	if err != nil {
+		log.Println(err)
+		return false
+	}
+
+	if reply.System.GetSysinfo.ErrCode != 0 {
+		log.Printf("Failed to set relay state, error: %d\n", reply.System.GetSysinfo.ErrCode)
+
+		return false
+	}
+
+	return reply.System.GetSysinfo.RelayState == 1
+}
+
diff --git a/integration/kasa/message.go b/integration/kasa/message.go
new file mode 100644
index 0000000..2b2b71f
--- /dev/null
+++ b/integration/kasa/message.go
@@ -0,0 +1,26 @@
+package kasa
+
+type reply struct {
+	System struct {
+		SetRelayState struct {
+			ErrCode int `json:"err_code"`
+		} `json:"set_relay_state"`
+		GetSysinfo struct {
+			RelayState int `json:"relay_state"`
+			ErrCode int `json:"err_code"`
+		} `json:"get_sysinfo"`
+	} `json:"system"`
+}
+
+type SetRelayState struct {
+	State int `json:"state"`
+}
+
+type GetSysinfo struct{}
+
+type cmd struct {
+	System struct {
+		SetRelayState *SetRelayState `json:"set_relay_state,omitempty"`
+		GetSysinfo    *GetSysinfo    `json:"get_sysinfo,omitempty"`
+	} `json:"system"`
+}
diff --git a/integration/kasa/reply.go b/integration/kasa/reply.go
deleted file mode 100644
index fc279ec..0000000
--- a/integration/kasa/reply.go
+++ /dev/null
@@ -1,19 +0,0 @@
-package kasa
-
-type errCode struct {
-	ErrCode int
-}
-
-type reply struct {
-	System struct {
-		SetRelayState errCode `json:"set_relay_state"`
-	} `json:"system"`
-}
-
-type cmd struct {
-	System struct {
-		SetRelayState struct {
-			State int `json:"state"`
-		} `json:"set_relay_state"`
-	} `json:"system"`
-}