package protocol

import (
	"bufio"
	"bytes"
	"crypto/tls"
	_ "embed"
	"fmt"
	"io"
	"mime/multipart"
	"net"
	"net/http"
	"net/textproto"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/vulncheck-oss/go-exploit/db"
	"github.com/vulncheck-oss/go-exploit/output"
	"github.com/vulncheck-oss/go-exploit/transform"
)

// GlobalUA is the default User-Agent for all go-exploit comms
//
//go:embed http-user-agent.txt
var GlobalUA string

// GlobalCommTimeout is the default timeout for all socket communications.
var GlobalCommTimeout = 10

// Returns a valid HTTP/HTTPS URL provided the given input.
func GenerateURL(rhost string, rport int, ssl bool, uri string) string {
	url := ""
	if ssl {
		url += "https://"
	} else {
		url += "http://"
	}

	// is the address v6?
	ip := net.ParseIP(rhost)
	if ip != nil && ip.To4() == nil {
		rhost = "[" + rhost + "]"
	}

	url += rhost
	url += ":"
	url += strconv.Itoa(rport)
	url += uri

	return url
}

// Using the variable amount of paths, return a URI without any extra '/'.
func BuildURI(paths ...string) string {
	uri := "/"
	for _, path := range paths {
		if !strings.HasSuffix(uri, "/") && !strings.HasPrefix(path, "/") {
			uri += "/"
		}
		uri += path
	}

	return uri
}

// BasicAuth takes a username and password and returns a string suitable for an Authorization header.
func BasicAuth(username, password string) string {
	return "Basic " + transform.EncodeBase64(username+":"+password)
}

func parseCookies(headers []string) string {
	cookies := make([]string, len(headers))

	for i, cookie := range headers {
		cookies[i] = strings.Split(cookie, ";")[0]
	}

	return strings.Join(cookies, "; ")
}

// ParseCookies parses an HTTP response and returns a string suitable for a Cookie header.
func ParseCookies(resp *http.Response) string {
	return parseCookies(resp.Header.Values("Set-Cookie"))
}

// Go doesn't always like sending our exploit URI so use this raw version. SSL not implemented.
func DoRawHTTPRequest(rhost string, rport int, uri string, verb string) bool {
	// connect
	conn, success := TCPConnect(rhost, rport)
	if !success {
		return false
	}

	// is the address v6?
	ip := net.ParseIP(rhost)
	if ip != nil && ip.To4() == nil {
		rhost = "[" + rhost + "]"
	}

	httpRequest := verb + " " + uri + " HTTP/1.1\r\n"
	httpRequest += "Host: " + rhost + ":" + strconv.Itoa(rport) + "\r\n"
	if len(GlobalUA) != 0 {
		httpRequest += "User-Agent: " + GlobalUA + "\r\n"
	}
	httpRequest += "Accept: */*\r\n"
	httpRequest += "\r\n"
	success = TCPWrite(conn, []byte(httpRequest))
	if !success {
		return false
	}

	// don't currently care about the response. Read a byte and move on'
	_, success = TCPReadAmount(conn, 1)

	return success
}

// Cache the respsone in the database for later reuse.
func cacheResponse(req *http.Request, resp *http.Response) {
	parsedURL, _ := url.Parse(req.URL.String())
	port, err := strconv.Atoi(parsedURL.Port())
	if err != nil {
		output.PrintFrameworkError(err.Error())

		return
	}

	respBuffer := &bytes.Buffer{}
	err = resp.Write(respBuffer)
	if err != nil {
		output.PrintfFrameworkError("Resp write error: %s", err.Error())

		return
	}

	db.CacheHTTPResponse(parsedURL.Hostname(), port, parsedURL.Path, respBuffer.Bytes())
}

// Look up matching URI in the HTTP cache and return it if found.
func cacheLookup(uri string) (*http.Response, string, bool) {
	parsedURL, _ := url.Parse(uri)
	port, err := strconv.Atoi(parsedURL.Port())
	if err != nil {
		output.PrintFrameworkError(err.Error())

		return nil, "", false
	}

	cachedResp, ok := db.GetHTTPResponse(parsedURL.Hostname(), port, parsedURL.Path)
	if !ok {
		// didn't get any cache data. no big deal.
		return nil, "", false
	}

	resp, err := http.ReadResponse(bufio.NewReader(strings.NewReader(cachedResp)), nil)
	if err != nil {
		output.PrintFrameworkError(err.Error())

		return nil, "", false
	}
	defer resp.Body.Close()

	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		// seen this fail when, for example, Shodan messes with chunking
		output.PrintFrameworkError(err.Error())

		return nil, "", false
	}
	if bytes.HasPrefix(bodyBytes, []byte("\x1f\x8b\x08")) {
		// if the data in the cache is still compressed, decompress it
		bodyBytes, ok = transform.Inflate(bodyBytes)
		if !ok {
			return nil, "", false
		}
	}

	output.PrintfFrameworkTrace("HTTP cache hit: %s", uri)

	bodyString := string(bodyBytes)

	return resp, bodyString, true
}

// Provided an HTTP client and a req, this function triggers the HTTP request and converts
// the response body to a string.
func DoRequest(client *http.Client, req *http.Request) (*http.Response, string, bool) {
	resp, err := client.Do(req)
	if err != nil {
		output.PrintfFrameworkError("HTTP request error: %s", err)

		return resp, "", false
	}
	defer resp.Body.Close()

	bodyBytes, _ := io.ReadAll(resp.Body)

	return resp, string(bodyBytes), true
}

// Turns `net/http` []*Cookie into a string for adding to the Cookie header.
func CookieString(cookies []*http.Cookie) string {
	cookieString := ""
	for c, cookie := range cookies {
		if c == 0 {
			cookieString += cookie.Name + "=" + cookie.Value + ";"
		} else {
			cookieString += " " + cookie.Name + "=" + cookie.Value + ";"
		}
	}

	return cookieString
}

// converts a map of strings into a single string in application/x-www-urlencoded format (but does not encode the params!)
func CreateRequestParams(params map[string]string) string {
	data := ""
	for key, element := range params {
		if len(data) > 0 {
			data += "&"
		}
		data += (key + "=" + element)
	}

	return data
}

// CreateRequestParamsEncoded is the encoded version of CreateRequestParams.
func CreateRequestParamsEncoded(params map[string]string) string {
	paramsCopy := make(map[string]string)

	for k, v := range params {
		paramsCopy[k] = url.QueryEscape(v)
	}

	return CreateRequestParams(paramsCopy)
}

// Provided a map of headers, this function loops through them and sets them in the http request.
func SetRequestHeaders(req *http.Request, headers map[string]string) {
	for key, value := range headers {
		if key == "Host" {
			// host can't be set directly
			req.Host = value
		} else {
			// don't use the Set function because the module might modify key. Set the header directly.
			req.Header[key] = []string{value}
		}
	}
}

// Creates the HTTP client, generates the HTTP request, and sets the default user-agent.
func CreateRequest(verb string, url string, payload string, followRedirect bool) (*http.Client, *http.Request, bool) {
	var client *http.Client
	if !followRedirect {
		client = &http.Client{
			Transport: &http.Transport{
				Proxy: http.ProxyFromEnvironment,
				Dial: (&net.Dialer{
					Timeout: time.Duration(GlobalCommTimeout) * time.Second,
				}).Dial,
				TLSClientConfig: (&tls.Config{
					InsecureSkipVerify: true,
					// We have no control over the SSL versions supported on the remote target. Be permissive for more targets.
					MinVersion: tls.VersionSSL30,
				}),
			},
			Timeout: time.Duration(GlobalCommTimeout) * time.Second,
			CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
				return http.ErrUseLastResponse
			},
		}
	} else {
		client = &http.Client{
			Transport: &http.Transport{
				Proxy: http.ProxyFromEnvironment,
				Dial: (&net.Dialer{
					Timeout: time.Duration(GlobalCommTimeout) * time.Second,
				}).Dial,
				TLSClientConfig: (&tls.Config{
					InsecureSkipVerify: true,
					// We have no control over the SSL versions supported on the remote target. Be permissive for more targets.
					MinVersion: tls.VersionSSL30,
				}),
			},
			Timeout: time.Duration(GlobalCommTimeout) * time.Second,
		}
	}

	req, err := http.NewRequest(verb, url, strings.NewReader(payload))
	if err != nil {
		output.PrintfFrameworkError("HTTP request creation error: %s", err)

		return nil, nil, false
	}

	// set headers on the request
	req.Header.Set("User-Agent", GlobalUA)

	return client, req, true
}

// Generic send HTTP request and receive response.
func HTTPSendAndRecv(verb string, url string, payload string) (*http.Response, string, bool) {
	client, req, ok := CreateRequest(verb, url, payload, true)
	if !ok {
		return nil, "", false
	}

	return DoRequest(client, req)
}

func HTTPGetCache(url string) (*http.Response, string, bool) {
	// first see if we have it cached somewhere
	if db.GlobalSQLHandle != nil {
		resp, body, ok := cacheLookup(url)
		if ok {
			return resp, body, true
		}
	}

	client, req, ok := CreateRequest("GET", url, "", true)
	if !ok {
		return nil, "", false
	}

	resp, err := client.Do(req)
	if err != nil {
		output.PrintfFrameworkError("HTTP request error: %s", err)

		return resp, "", false
	}
	defer resp.Body.Close()

	bodyBytes, _ := io.ReadAll(resp.Body)
	bodyString := string(bodyBytes)

	if db.GlobalSQLHandle != nil {
		// shove the body back in to be re-read for storage
		resp.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
		cacheResponse(req, resp)
	}

	return resp, bodyString, true
}

// Send an HTTP request but do not follow the 302 redirect.
func HTTPSendAndRecvNoRedirect(verb string, url string, payload string) (*http.Response, string, bool) {
	client, req, ok := CreateRequest(verb, url, payload, true)
	if !ok {
		return nil, "", false
	}

	// ignore the redirect
	client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
		return http.ErrUseLastResponse
	}

	return DoRequest(client, req)
}

// Send an HTTP request, with the provided parameters in the params map stored in the body.
// Return the response and response body.
//
// Note that this function *will not* attempt to url encode the params.
func HTTPSendAndRecvURLEncoded(verb string, url string, params map[string]string) (*http.Response, string, bool) {
	payload := CreateRequestParams(params)
	client, req, ok := CreateRequest(verb, url, payload, true)
	if !ok {
		return nil, "", false
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	return DoRequest(client, req)
}

// Send an HTTP request, with the provided parameters in the params map URL encoded in the body.
// Return the response and response body.
//
// Note that this function *will* attempt to url encode the params.
func HTTPSendAndRecvURLEncodedParams(verb string, url string, params map[string]string) (*http.Response, string, bool) {
	payload := CreateRequestParamsEncoded(params)
	client, req, ok := CreateRequest(verb, url, payload, true)
	if !ok {
		return nil, "", false
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	return DoRequest(client, req)
}

// Send an HTTP request, with the provided parameters in the params map stored in the body, and
// with extra headers specified in the headers map. Return the response and response body.
//
// Note that this function *will not* attempt to url encode the params.
func HTTPSendAndRecvURLEncodedAndHeaders(verb string, url string, params map[string]string,
	headers map[string]string,
) (*http.Response, string, bool) {
	payload := CreateRequestParams(params)

	client, req, ok := CreateRequest(verb, url, payload, true)
	if !ok {
		return nil, "", false
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	SetRequestHeaders(req, headers)

	return DoRequest(client, req)
}

// Send an HTTP request, with the provided parameters in the params map URL encoded in the body, and
// with extra headers specified in the headers map. Return the response and response body.
//
// Note that this function *will* attempt to url encode the params.
func HTTPSendAndRecvURLEncodedParamsAndHeaders(verb string, url string, params map[string]string,
	headers map[string]string,
) (*http.Response, string, bool) {
	payload := CreateRequestParamsEncoded(params)

	client, req, ok := CreateRequest(verb, url, payload, true)
	if !ok {
		return nil, "", false
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	SetRequestHeaders(req, headers)

	return DoRequest(client, req)
}

// Send an HTTP request with extra headers specified in the headers map. Return the response and response body.
func HTTPSendAndRecvWithHeaders(verb string, url string, payload string, headers map[string]string) (*http.Response, string, bool) {
	client, req, ok := CreateRequest(verb, url, payload, true)
	if !ok {
		return nil, "", false
	}

	SetRequestHeaders(req, headers)

	return DoRequest(client, req)
}

// this naming scheme is a little out of control.
func HTTPSendAndRecvWithHeadersNoRedirect(verb string, url string, payload string,
	headers map[string]string,
) (*http.Response, string, bool) {
	client, req, ok := CreateRequest(verb, url, payload, true)
	if !ok {
		return nil, "", false
	}

	// ignore the redirect
	client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
		return http.ErrUseLastResponse
	}

	SetRequestHeaders(req, headers)

	return DoRequest(client, req)
}

func MultipartCreateForm() (*strings.Builder, *multipart.Writer) {
	form := &strings.Builder{}
	w := multipart.NewWriter(form)

	return form, w
}

func MultipartAddField(writer *multipart.Writer, name string, value string) bool {
	fw, err := writer.CreateFormField(name)
	if err != nil {
		return false
	}
	_, err = io.Copy(fw, strings.NewReader(value))

	return err == nil
}

// MultipartCreateFormFields generates multipart form data out of the field names and values
// provided via fieldMap and writes this to writer.
// It returns a bool to indicate success or failure.
func MultipartCreateFormFields(writer *multipart.Writer, fieldMap map[string]string) bool {
	for fieldName, value := range fieldMap {
		if ok := MultipartAddField(writer, fieldName, value); !ok {
			return false
		}
	}

	return true
}

func MultipartAddPart(writer *multipart.Writer, headers map[string]string, body string) bool {
	h := make(textproto.MIMEHeader)
	for k, v := range headers {
		h.Set(k, v)
	}

	fw, err := writer.CreatePart(h)
	if err != nil {
		return false
	}
	_, err = io.Copy(fw, strings.NewReader(body))

	return err == nil
}

func MultipartAddFile(writer *multipart.Writer, name, filename, ctype, value string) bool {
	// CreateFormFile doesn't expose Content-Type
	return MultipartAddPart(writer, map[string]string{
		"Content-Disposition": fmt.Sprintf(`form-data; name="%s"; filename="%s"`, name, filename),
		"Content-Type":        ctype,
	}, value)
}

// Provided an HTTP request, find the Set-Cookie headers, and extract
// the value of the specified cookie. Example:.
func GetSetCookieValue(resp *http.Response, name string) (string, bool) {
	cookies, ok := resp.Header["Set-Cookie"]
	if !ok {
		output.PrintError("Missing Set-Cookie header")

		return "", false
	}

	for _, entry := range cookies {
		if strings.HasPrefix(entry, name+"=") {
			end := len(entry)
			index := strings.Index(entry, ";")
			if index != -1 {
				end = index
			}

			return entry[len(name+"="):end], true
		}
	}

	return "", false
}