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)

	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)                       // 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])
	pkt := make([]byte, pktlen+padlen)
	if _, err := io.ReadFull(c.Conn, pkt); err != nil {
		return 0, err
	}

	// ECB mode here too.........
	bs := c.block.BlockSize()
	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
}