/*
   Copyright 2023 Josh Deprez

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

	   http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
)

const (
	privateKey = "Grow#0*2Sun68CbE"                                                 // 16 bytes
	noCrypto1  = "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" // 16 bytes
	noCrypto2  = "\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff" // 16 bytes
	getKeyMsg  = "\x68\x68\x00\x00\x00\x06\xf7\x04\x0a\xe7\x00\x08"
	header     = "\x68\x68"
	padding    = noCrypto2 // all FF bytes
)

func getKey(conn io.ReadWriter) ([]byte, error) {
	if _, err := conn.Write([]byte(getKeyMsg)); err != nil {
		return nil, err
	}
	kp := make([]byte, 25)
	if _, err := io.ReadFull(conn, kp); err != nil {
		return nil, err
	}
	key := kp[9:] // must be 16 bytes long, because maths
	if bytes.Equal(key, []byte(noCrypto1)) || bytes.Equal(key, []byte(noCrypto2)) {
		return nil, nil
	}
	for i := range key {
		key[i] ^= privateKey[i]
	}
	return key, nil
}

// sungrowConn wraps a regular TCP connection with funky Sungrow encryption.
type sungrowConn struct {
	net.Conn
	block cipher.Block
	fifo  *bytes.Buffer

	mu   sync.RWMutex
	txid uint16
}

func (c *sungrowConn) Write(msg []byte) (int, error) {
	if c.block == nil {
		return c.Conn.Write(msg)
	}

	c.mu.Lock()
	defer c.mu.Unlock()
	c.txid = binary.BigEndian.Uint16(msg[:2])
	c.fifo.Truncate(0)

	bs := c.block.BlockSize()
	padlen := bs - (len(msg) % bs)
	if padlen == bs {
		padlen = 0
	}

	req := make([]byte, 4+len(msg)+padlen)
	copy(req, []byte{1, 0, byte(len(msg)), byte(padlen)}) // 4 byte encryption header
	copy(req[4:], header)                                 // 2 byte sungrow header, replaces txid
	copy(req[6:], msg[2:])                                // rest of message
	copy(req[4+len(msg):], padding[:padlen])              // padding

	// ECB mode.........
	for cp := req[4:]; len(cp) > 0; cp = cp[bs:] {
		c.block.Encrypt(cp, cp)
	}

	n, err := c.Conn.Write(req)
	return n - 4 - padlen, err
}

func (c *sungrowConn) Read(out []byte) (int, error) {
	if c.block == nil {
		return c.Conn.Read(out)
	}

	if c.fifo.Len() > 0 {
		c.mu.RLock()
		defer c.mu.RUnlock()
		return c.fifo.Read(out)
	}
	// 4-byte header describes how much to read
	hdr := make([]byte, 4)
	if _, err := io.ReadFull(c.Conn, hdr); err != nil {
		return 0, err
	}
	pktlen, padlen := int(hdr[2]), int(hdr[3])
	bs := c.block.BlockSize()
	if (pktlen+padlen)%bs != 0 {
		return 0, fmt.Errorf("pktlen + padlen = %d + %d = %d, want divisible by %d", pktlen, padlen, pktlen+padlen, bs)
	}
	pkt := make([]byte, pktlen+padlen)
	if _, err := io.ReadFull(c.Conn, pkt); err != nil {
		return 0, err
	}

	// ECB mode here too.........
	for cp := pkt; len(cp) > 0; cp = cp[bs:] {
		c.block.Decrypt(cp, cp)
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	binary.BigEndian.PutUint16(pkt, c.txid)
	c.fifo.Write(pkt[:pktlen])
	return c.fifo.Read(out)
}

// dialSungrow dials a TCP connection, obtains the encryption key, and returns
// a sungrowConn that encrypts/decrypts using the key.
func dialSungrow(addr string) (*sungrowConn, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, fmt.Errorf("dialing: %w", err)
	}
	key, err := getKey(conn)
	if err != nil {
		return nil, fmt.Errorf("obtaining key: %w", err)
	}
	if len(key) == 0 { // no encryption
		return &sungrowConn{Conn: conn}, nil
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("setting cipher: %w", err)
	}
	return &sungrowConn{
		Conn:  conn,
		block: block,
		fifo:  new(bytes.Buffer),
	}, nil
}