package client

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
	"testing"
	"time"

	"teleport/pkg/config"
	"teleport/pkg/encryption"
)

func testClient(t *testing.T) *TeleportClient {
	t.Helper()
	return NewTeleportClient(&config.Config{
		EncryptionKey: "a0e3dd20a761b118ca234160dd8b87230a001e332a97c9cfe3b8b9c99efaae03",
	})
}

func TestForwardDataHidesHTTPPlaintextOnTunnel(t *testing.T) {
	tc := testClient(t)
	defer tc.Stop()

	localA, localB := net.Pipe()
	tunnelA, tunnelB := net.Pipe()
	defer localA.Close()
	defer localB.Close()
	defer tunnelA.Close()
	defer tunnelB.Close()

	go tc.forwardData(localA, tunnelA, true)

	req := []byte("GET /secret HTTP/1.1\r\nHost: example.com\r\n\r\n")
	writeErr := make(chan error, 1)
	go func() {
		_, err := localB.Write(req)
		writeErr <- err
	}()

	tunnelB.SetReadDeadline(time.Now().Add(5 * time.Second))
	var length uint32
	if err := binary.Read(tunnelB, binary.BigEndian, &length); err != nil {
		t.Fatalf("failed to read frame length from tunnel: %v", err)
	}
	if length == 0 || length > encryption.MaxFrameSize {
		t.Fatalf("unexpected frame length %d", length)
	}
	ciphertext := make([]byte, length)
	if _, err := io.ReadFull(tunnelB, ciphertext); err != nil {
		t.Fatalf("failed to read ciphertext from tunnel: %v", err)
	}

	if bytes.Contains(ciphertext, []byte("GET /")) ||
		bytes.Contains(ciphertext, []byte("Host:")) ||
		bytes.Contains(ciphertext, []byte("example.com")) {
		t.Fatal("HTTP plaintext visible on client->server tunnel")
	}

	var frame bytes.Buffer
	if err := binary.Write(&frame, binary.BigEndian, length); err != nil {
		t.Fatal(err)
	}
	frame.Write(ciphertext)
	plain, err := encryption.ReadEncryptedFrame(&frame, tc.derivedKey)
	if err != nil {
		t.Fatalf("decrypt of tunneled frame failed: %v", err)
	}
	if !bytes.Equal(plain, req) {
		t.Errorf("decrypted payload mismatch: got %q want %q", plain, req)
	}

	select {
	case err := <-writeErr:
		if err != nil {
			t.Fatalf("write to local conn failed: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for local write")
	}
}

func TestForwardDataDecryptsTunnelToLocal(t *testing.T) {
	tc := testClient(t)
	defer tc.Stop()

	tunnelA, tunnelB := net.Pipe()
	localA, localB := net.Pipe()
	defer tunnelA.Close()
	defer tunnelB.Close()
	defer localA.Close()
	defer localB.Close()

	go tc.forwardData(tunnelA, localA, false)

	payload := []byte("HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nhello")
	writeErr := make(chan error, 1)
	go func() {
		writeErr <- encryption.WriteEncryptedFrame(tunnelB, payload, tc.derivedKey)
	}()

	localB.SetReadDeadline(time.Now().Add(5 * time.Second))
	got := make([]byte, len(payload))
	if _, err := io.ReadFull(localB, got); err != nil {
		t.Fatalf("failed to read decrypted payload: %v", err)
	}
	if !bytes.Equal(got, payload) {
		t.Errorf("got %q want %q", got, payload)
	}

	select {
	case err := <-writeErr:
		if err != nil {
			t.Fatalf("WriteEncryptedFrame failed: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for tunnel write")
	}
}