package cli

import (
	"strings"
	"testing"

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

func TestCodeExecutionCmdLineParse(t *testing.T) {
	conf := config.New(config.CodeExecution, []c2.Impl{c2.SimpleShellServer}, "test product", "CVE-2023-1270", 1270)
	conf.Rhost = "rcetest"

	success := CodeExecutionCmdLineParse(conf)

	if conf.Rhost != "" {
		t.Fatal("Rhost should have no default value")
	}
	if conf.Rport != 1270 {
		t.Fatal("Rport should default to passed in value")
	}
	if conf.SSL != false {
		t.Fatal("SSL should default to false")
	}
	if conf.DoVerify != false {
		t.Fatal("verify should default to false")
	}
	if conf.DoVersionCheck != false {
		t.Fatal("version check should default to false")
	}
	if conf.DoExploit != false {
		t.Fatal("exploit should default to false")
	}
	if success != false {
		t.Fatal("parsing should have failed")
	}
	if conf.ThirdPartyC2Server != false {
		t.Fatal("outside should default to false")
	}
	if conf.C2Timeout != 30 {
		t.Fatal("timeout should default to 30")
	}
}

func TestCommonValidate(t *testing.T) {
	conf := config.New(config.CodeExecution, []c2.Impl{c2.SimpleShellServer}, "test product", "CVE-2023-1270", 1270)
	var rhosts string
	var rports string
	var rhostsFile string

	if commonValidate(conf, rhosts, rports, rhostsFile) {
		t.Fatal("commonValidate should fail with an empty Rhost")
	}

	conf.Rhost = "10.9.49.99"
	if commonValidate(conf, rhosts, rports, rhostsFile) {
		t.Fatal("commonValidate should fail with no supplied action")
	}

	conf.DoVerify = true
	if !commonValidate(conf, rhosts, rports, rhostsFile) {
		t.Fatal("commonValidate should succeed with rhost, rport, and doVerify")
	}

	conf.Rhost = ""
	if !commonValidate(conf, "127.0.0.1", "1270,1280", rhostsFile) {
		t.Fatal("commonValidate should have succeeded")
	}

	if !commonValidate(conf, "127.0.0.1,127.0.0.2", rports, rhostsFile) {
		t.Fatal("commonValidate have succeeded")
	}
}

func TestRhostsParsing(t *testing.T) {
	conf := config.New(config.CodeExecution, []c2.Impl{c2.SimpleShellServer}, "test product", "CVE-2023-1270", 1270)

	if !handleRhostsOptions(conf, "127.0.0.1,127.0.0.2", "80,443", "") {
		t.Fatal("commonValidate should succeed")
	}
	if len(conf.RhostsNTuple) != 4 {
		t.Fatal("Failed to parse rhosts")
	}
	if conf.RhostsNTuple[0].Rhost != "127.0.0.1" || conf.RhostsNTuple[1].Rhost != "127.0.0.1" ||
		conf.RhostsNTuple[2].Rhost != "127.0.0.2" || conf.RhostsNTuple[3].Rhost != "127.0.0.2" {
		t.Fatal("Failed to parse rhosts")
	}
	if conf.RhostsNTuple[0].Rport != 80 || conf.RhostsNTuple[1].Rport != 443 {
		t.Fatal("Failed to parse rports")
	}
	conf.RhostsNTuple = make([]config.RhostTriplet, 0)

	if !handleRhostsOptions(conf, "127.0.0.3", "443", "") {
		t.Fatal("commonValidate should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Failed to parse rhosts")
	}
	if conf.RhostsNTuple[0].Rhost != "127.0.0.3" {
		t.Fatal("Failed to parse rhosts")
	}
	if conf.RhostsNTuple[0].Rport != 443 {
		t.Fatal("Failed to parse rports")
	}
	conf.RhostsNTuple = make([]config.RhostTriplet, 0)

	conf.Rhost = "127.0.0.4"
	if !handleRhostsOptions(conf, "", "443,80,8080", "") {
		t.Fatal("commonValidate should succeed")
	}
	if len(conf.RhostsNTuple) != 3 {
		t.Fatal("Failed to parse rhosts")
	}
	if conf.RhostsNTuple[0].Rhost != "127.0.0.4" {
		t.Fatal("Failed to parse rhosts")
	}
	if conf.RhostsNTuple[0].Rport != 443 {
		t.Fatal("Failed to parse rports")
	}
	if conf.RhostsNTuple[1].Rport != 80 {
		t.Fatal("Failed to parse rports")
	}
	if conf.RhostsNTuple[2].Rport != 8080 {
		t.Fatal("Failed to parse rports")
	}
	conf.Rhost = ""

	conf.RhostsNTuple = make([]config.RhostTriplet, 0)
	if !handleRhostsOptions(conf, "192.168.1.0/24", "80", "") {
		t.Fatal("commonValidate should succeed")
	}
	if len(conf.RhostsNTuple) != 256 {
		t.Fatal("Failed to parse rhosts")
	}

	conf.RhostsNTuple = make([]config.RhostTriplet, 0)
	if !handleRhostsOptions(conf, "192.168.1.0/24", "80,8080", "") {
		t.Fatal("commonValidate should succeed")
	}
	if len(conf.RhostsNTuple) != 512 {
		t.Fatal("Failed to parse rhosts")
	}
}

func TestParseTuple(t *testing.T) {
	conf := config.New(config.CodeExecution, []c2.Impl{c2.SimpleShellServer}, "test product", "CVE-2023-1270", 1270)

	if !parseTuple(conf, "example.com:80") {
		t.Fatal("ParseTuple should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Missing Rhosts entry")
	}
	if conf.RhostsNTuple[0].Rport != 80 {
		t.Fatal("Invalid tuple parsing")
	}
	if conf.RhostsNTuple[0].Rhost != "example.com" {
		t.Fatal("Invalid tuple parsing")
	}

	conf.RhostsNTuple = make([]config.RhostTriplet, 0)
	if parseTuple(conf, "example.com") {
		t.Fatal("ParseTuple should have failed")
	}
	if parseTuple(conf, "example.com:hi") {
		t.Fatal("ParseTuple should have failed")
	}

	if !parseTuple(conf, "127.0.0.5:8080") {
		t.Fatal("ParseTuple should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Missing Rhosts entry")
	}
	if conf.RhostsNTuple[0].Rport != 8080 {
		t.Fatal("Invalid tuple parsing")
	}
	if conf.RhostsNTuple[0].Rhost != "127.0.0.5" {
		t.Fatal("Invalid tuple parsing")
	}

	conf.RhostsNTuple = make([]config.RhostTriplet, 0)
	if !parseTuple(conf, "fe80::1b8c:b574:ebec:81c0:65531") {
		t.Fatal("ParseTuple should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Missing Rhosts entry")
	}
	if conf.RhostsNTuple[0].Rport != 65531 {
		t.Fatal("Invalid tuple parsing")
	}
	if conf.RhostsNTuple[0].Rhost != "fe80::1b8c:b574:ebec:81c0" {
		t.Fatal("Invalid tuple parsing")
	}
}

func TestParseTriplet(t *testing.T) {
	conf := config.New(config.CodeExecution, []c2.Impl{c2.SimpleShellServer}, "test product", "CVE-2023-1270", 1270)

	if !parseTriplet(conf, "example.com,80,") {
		t.Fatal("parseTriplet should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Missing Rhosts entry")
	}
	if conf.RhostsNTuple[0].Rport != 80 {
		t.Fatal("Invalid triplet parsing")
	}
	if conf.RhostsNTuple[0].Rhost != "example.com" {
		t.Fatal("Invalid triplet parsing")
	}

	conf.RhostsNTuple = make([]config.RhostTriplet, 0)
	if parseTriplet(conf, "example.com,,") {
		t.Fatal("parseTriplet should have failed")
	}
	if parseTriplet(conf, "example.commhi,") {
		t.Fatal("parseTriplet should have failed")
	}

	if !parseTriplet(conf, "127.0.0.3,8080,") {
		t.Fatal("parseTriplet should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Missing Rhosts entry")
	}
	if conf.RhostsNTuple[0].Rport != 8080 {
		t.Fatal("Invalid triplet parsing")
	}
	if conf.RhostsNTuple[0].Rhost != "127.0.0.3" {
		t.Fatal("Invalid triplet parsing")
	}

	conf.RhostsNTuple = make([]config.RhostTriplet, 0)
	if !parseTriplet(conf, "fe80::1b8c:b574:ebec:81c0,65531,") {
		t.Fatal("parseTriplet should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Missing Rhosts entry")
	}
	if conf.RhostsNTuple[0].Rport != 65531 {
		t.Fatal("Invalid triplet parsing")
	}
	if conf.RhostsNTuple[0].Rhost != "fe80::1b8c:b574:ebec:81c0" {
		t.Fatal("Invalid triplet parsing")
	}
}

func TestParseIntelJSON(t *testing.T) {
	conf := config.New(config.CodeExecution, []c2.Impl{c2.SimpleShellServer}, "test product", "CVE-2023-1270", 1270)

	intel := `{"ip":"10.9.49.2","port":80,"ssl":false,"lastSeen":"2024-02-18T23:56:34.454071","asn":"AS33915","country":"Netherlands","country_code":"NL","city":"Hoofddorp","cve":["CVE-2021-36260"],"matches":["Hikvision IP Camera Web Language Command Injection"],"hostnames":["89-220-147-28.cable.dynamic.v4.ziggo.nl"],"type":{"id":"initial-access","finding":"potentially vulnerable"},"feed_ids":["a65d5009-f84b-4c3d-8803-1f8b1246ddeb"]}`

	if !parseIPIntel(conf, intel) {
		t.Fatal("parseTriplet should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Missing Rhosts entry")
	}
	if conf.RhostsNTuple[0].Rport != 80 {
		t.Fatal("Invalid triplet parsing")
	}
	if conf.RhostsNTuple[0].Rhost != "10.9.49.2" {
		t.Fatal("Invalid triplet parsing")
	}
	if conf.RhostsNTuple[0].SSL != config.SSLDisabled {
		t.Fatal("Invalid triplet parsing")
	}

	conf.RhostsNTuple = make([]config.RhostTriplet, 0)
	intel = `{"ip":"fe80::1b8c:b574:ebec:81c8","port":81,"ssl":true,"lastSeen":"2024-02-18T23:56:34.454071","asn":"AS33915","country":"Netherlands","country_code":"NL","city":"Hoofddorp","cve":["CVE-2021-36260"],"matches":["Hikvision IP Camera Web Language Command Injection"],"hostnames":["89-220-147-28.cable.dynamic.v4.ziggo.nl"],"type":{"id":"initial-access","finding":"potentially vulnerable"},"feed_ids":["a65d5009-f84b-4c3d-8803-1f8b1246ddeb"]}`
	if !parseIPIntel(conf, intel) {
		t.Fatal("parseTriplet should succeed")
	}
	if len(conf.RhostsNTuple) != 1 {
		t.Fatal("Missing Rhosts entry")
	}
	if conf.RhostsNTuple[0].Rport != 81 {
		t.Fatal("Invalid triplet parsing")
	}
	if conf.RhostsNTuple[0].Rhost != "fe80::1b8c:b574:ebec:81c8" {
		t.Fatal("Invalid triplet parsing")
	}
	if conf.RhostsNTuple[0].SSL != config.SSLEnabled {
		t.Fatal("Invalid triplet parsing")
	}
}

func TestFileFormatParams(t *testing.T) {
	conf := config.NewLocal(config.FileFormat, []c2.Impl{}, "test product", "CVE-2023-1270")
	if !loadFileFormatTemplate("", conf) {
		t.Fatal("an empty file path should pass validation")
	}
	if !loadFileFormatTemplate("commandline_test.go", conf) {
		t.Fatal("failed to load the file format template")
	}
	if !strings.Contains(conf.FileTemplateData, "TestFileFormatParams") {
		t.Fatal("failed to load the file format template")
	}
}