package tcp

import (
	"testing"

	"git.difuse.io/Difuse/Mellaris/analyzer"
)

func TestSSHAnalyzer_Name(t *testing.T) {
	a := &SSHAnalyzer{}
	if a.Name() != "ssh" {
		t.Errorf("Name() = %q, want ssh", a.Name())
	}
}

func TestSSHStream_Feed_Client(t *testing.T) {
	s := newSSHStream(nil)
	u, done := s.Feed(false, false, false, 0, []byte("SSH-2.0-OpenSSH_8.9p1 Ubuntu-3\r\n"))
	if u == nil {
		t.Fatal("Feed returned nil update")
	}
	if done {
		t.Error("should not be done (server not yet received)")
	}
	client, ok := u.M["client"].(analyzer.PropMap)
	if !ok {
		t.Fatal("client prop missing")
	}
	if client["protocol"] != "2.0" {
		t.Errorf("protocol = %v, want 2.0", client["protocol"])
	}
	if client["software"] != "OpenSSH_8.9p1" {
		t.Errorf("software = %v, want OpenSSH_8.9p1", client["software"])
	}
	if comments, ok := client["comments"]; ok {
		if comments != "Ubuntu-3" {
			t.Errorf("comments = %v, want Ubuntu-3", comments)
		}
	}
}

func TestSSHStream_Feed_ClientWithComments(t *testing.T) {
	s := newSSHStream(nil)
	u, done := s.Feed(false, false, false, 0, []byte("SSH-2.0-OpenSSH_7.4 Ubuntu-3\r\n"))
	if u == nil {
		t.Fatal("Feed returned nil update")
	}
	if done {
		t.Error("should not be done")
	}
	client := u.M["client"].(analyzer.PropMap)
	if client["comments"] != "Ubuntu-3" {
		t.Errorf("comments = %v, want Ubuntu-3", client["comments"])
	}
}

func TestSSHStream_Feed_Both(t *testing.T) {
	s := newSSHStream(nil)
	s.Feed(false, false, false, 0, []byte("SSH-2.0-OpenSSH_8.9\r\n"))
	u, done := s.Feed(true, false, false, 0, []byte("SSH-2.0-dropbear_2022.83\r\n"))
	if u == nil {
		t.Fatal("Feed returned nil update")
	}
	if !done {
		t.Error("should be done after both sides")
	}
	server, ok := u.M["server"].(analyzer.PropMap)
	if !ok {
		t.Fatal("server prop missing")
	}
	if server["software"] != "dropbear_2022.83" {
		t.Errorf("server software = %v", server["software"])
	}
}

func TestSSHStream_Feed_NotSSH(t *testing.T) {
	s := newSSHStream(nil)
	u, done := s.Feed(false, false, false, 0, []byte("HTTP/1.1 200 OK\r\n"))
	if u != nil {
		t.Error("should return nil for non-SSH")
	}
	if !done {
		t.Error("should be cancelled (done) for non-SSH")
	}
}

func TestSSHStream_Feed_InvalidLine(t *testing.T) {
	s := newSSHStream(nil)
	u, done := s.Feed(false, false, false, 0, []byte("SSH-2.0-foo bar baz\r\n"))
	if u != nil {
		t.Error("should return nil for invalid line (>2 fields)")
	}
	if !done {
		t.Error("should be cancelled for invalid SSH line")
	}
}

func TestSSHStream_Feed_NoLineEnd(t *testing.T) {
	s := newSSHStream(nil)
	u, done := s.Feed(false, false, false, 0, []byte("SSH-2.0-OpenSSH"))
	if u != nil {
		t.Error("should return nil when no EOL found yet")
	}
	if done {
		t.Error("should not be done, waiting for more data")
	}
}

func TestSSHStream_Feed_IncompleteThenComplete(t *testing.T) {
	s := newSSHStream(nil)
	u, _ := s.Feed(false, false, false, 0, []byte("SSH-2.0-"))
	if u != nil {
		t.Error("first partial feed should return nil")
	}
	u, done := s.Feed(false, false, false, 0, []byte("Dropbear\r\n"))
	if u == nil {
		t.Fatal("second feed should return update")
	}
	if done {
		t.Error("should not be done (only client)")
	}
	client := u.M["client"].(analyzer.PropMap)
	if client["software"] != "Dropbear" {
		t.Errorf("software = %v", client["software"])
	}
}

func TestSSHStream_Feed_Skip(t *testing.T) {
	s := newSSHStream(nil)
	_, done := s.Feed(false, false, false, 5, []byte("data"))
	if !done {
		t.Error("skip != 0 should return done=true")
	}
}

func TestSSHStream_Feed_Empty(t *testing.T) {
	s := newSSHStream(nil)
	u, done := s.Feed(false, false, false, 0, []byte{})
	if u != nil || done {
		t.Error("empty data should return nil, false")
	}
}

func TestSSHStream_Close(t *testing.T) {
	s := newSSHStream(nil)
	s.clientBuf.Append([]byte("data"))
	s.serverBuf.Append([]byte("data"))
	s.clientMap = analyzer.PropMap{"key": "val"}
	u := s.Close(false)
	if u != nil {
		t.Error("Close should return nil")
	}
	if s.clientBuf.Len() != 0 || s.serverBuf.Len() != 0 {
		t.Error("Close should reset buffers")
	}
	if s.clientMap != nil || s.serverMap != nil {
		t.Error("Close should nil maps")
	}
}