package c2

import (
	"testing"

	"github.com/vulncheck-oss/go-exploit/c2/channel"
	"github.com/vulncheck-oss/go-exploit/c2/httpservefile"
)

func TestHTTPServeFileInit(t *testing.T) {
	impl, success := GetInstance(HTTPServeFile)
	if !success {
		t.Fatal("Failed to create HTTPServeFile")
	}

	if len(httpservefile.GetInstance().HostedFiles) != 0 {
		t.Fatal("Instance has a filename already")
	}
	httpservefile.GetInstance().AddFile("name", "random", []byte("data"))
	if len(httpservefile.GetInstance().HostedFiles) != 1 {
		t.Fatalf("Should have added a file: %d", len(httpservefile.GetInstance().HostedFiles))
	}

	httpservefile.GetInstance().FilesToServe = "factory.go"
	success = impl.Init(&channel.Channel{IPAddr: "127.0.0.2", Port: 1271, HTTPAddr: "127.0.0.1", HTTPPort: 1270, IsClient: true})
	if success {
		t.Fatal("Failed to check if it was invoked as a client")
	}

	success = impl.Init(&channel.Channel{IPAddr: "127.0.0.2", Port: 1271, HTTPAddr: "127.0.0.1", HTTPPort: 1270, IsClient: false})
	if !success {
		t.Fatal("Failed to successfully process well-formed init call")
	}

	// random name should have been generated
	if len(httpservefile.GetInstance().HostedFiles["factory.go"].RandomName) == 0 {
		t.Fatal("Instance did not generate a random filename")
	}
	if httpservefile.GetInstance().HTTPAddr != "127.0.0.1" {
		t.Fatal("Instance did not take up the http bind addr")
	}
	if httpservefile.GetInstance().HTTPPort != 1270 {
		t.Fatal("Instance did not take up the http bind port")
	}
}

func TestExternalModuleSingleton(t *testing.T) {
	_, success := GetInstance(Impl{Name: "", Category: InvalidCategory})
	if success {
		t.Fatal("Invalid & default External modules should not allow instance creation")
	}
	_, success = GetInstance(Impl{Name: "", Category: ExternalCategory})
	if success {
		t.Fatal("Empty names should not be allowed")
	}
	_, success = GetInstance(Impl{Name: "i1", Category: ExternalCategory})
	if !success {
		t.Fatal("Initial instance of external module was not returned")
	}
	_, success = GetInstance(Impl{Name: "i2", Category: ExternalCategory})
	if !success {
		t.Fatal("Could not get second ExternalCategory instance")
	}
}

func TestExternalModuleAddC2(t *testing.T) {
	_, ok := StringToImpl("blorp")
	if ok {
		t.Fatal("StringToImpl should not have returned a value")
	}
	i1 := AddC2("blorp")
	i2, ok := StringToImpl("blorp")
	if !ok {
		t.Fatal("StringToImpl should have returned a value")
	}
	if i1 != i2 {
		t.Fatal("Added C2 does not match")
	}
}