nuclei/v2/internal/runner/nucleicloud/cloud.go

package nucleicloud

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	jsoniter "github.com/json-iterator/go"
	"github.com/pkg/errors"
	"github.com/projectdiscovery/nuclei/v2/pkg/output"
	"github.com/projectdiscovery/retryablehttp-go"
)

// Client is a client for result retrieval from nuclei-cloud API
type Client struct {
	baseURL    string
	apiKey     string
	httpclient *retryablehttp.Client
}

const (
	pollInterval   = 1 * time.Second
	defaultBaseURL = "http://webapp.localhost"
	resultSize     = 100
)

// New returns a nuclei-cloud API client
func New(baseURL, apiKey string) *Client {
	options := retryablehttp.DefaultOptionsSingle
	options.Timeout = 15 * time.Second
	client := retryablehttp.NewClient(options)

	baseAppURL := baseURL
	if baseAppURL == "" {
		baseAppURL = defaultBaseURL
	}
	return &Client{httpclient: client, baseURL: baseAppURL, apiKey: apiKey}
}

// AddScan adds a scan for templates and target to nuclei server
func (c *Client) AddScan(req *AddScanRequest) (string, error) {
	var buf bytes.Buffer
	if err := jsoniter.NewEncoder(&buf).Encode(req); err != nil {
		return "", errors.Wrap(err, "could not json encode scan request")
	}
	httpReq, err := retryablehttp.NewRequest(http.MethodPost, fmt.Sprintf("%s/scan", c.baseURL), bytes.NewReader(buf.Bytes()))
	if err != nil {
		return "", errors.Wrap(err, "could not make request")
	}
	httpReq.Header.Set("X-API-Key", c.apiKey)

	resp, err := c.httpclient.Do(httpReq)
	if err != nil {
		return "", errors.Wrap(err, "could not do add scan request")
	}
	if resp.StatusCode != 200 {
		data, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return "", errors.Errorf("could not do request %d: %s", resp.StatusCode, string(data))
	}
	var data map[string]string
	if err := jsoniter.NewDecoder(resp.Body).Decode(&data); err != nil {
		resp.Body.Close()
		return "", errors.Wrap(err, "could not decode resp")
	}
	resp.Body.Close()
	id := data["id"]
	return id, nil
}

// GetResults gets results from nuclei server for an ID
// until there are no more results left to retrieve.
func (c *Client) GetResults(ID string, callback func(*output.ResultEvent), checkProgress bool) error {
	lastID := int64(0)
	for {
		uri := fmt.Sprintf("%s/results?id=%s&from=%d&size=%d", c.baseURL, ID, lastID, resultSize)
		httpReq, err := retryablehttp.NewRequest(http.MethodGet, uri, nil)
		if err != nil {
			return errors.Wrap(err, "could not make request")
		}
		httpReq.Header.Set("X-API-Key", c.apiKey)

		resp, err := c.httpclient.Do(httpReq)
		if err != nil {
			return errors.Wrap(err, "could not do ger result request")
		}
		if resp.StatusCode != 200 {
			data, _ := io.ReadAll(resp.Body)
			resp.Body.Close()
			return errors.Errorf("could not do request %d: %s", resp.StatusCode, string(data))
		}
		var items GetResultsResponse
		if err := jsoniter.NewDecoder(resp.Body).Decode(&items); err != nil {
			resp.Body.Close()
			return errors.Wrap(err, "could not decode results")
		}
		resp.Body.Close()

		for _, item := range items.Items {
			lastID = item.ID

			var result output.ResultEvent
			if err := jsoniter.NewDecoder(strings.NewReader(item.Raw)).Decode(&result); err != nil {
				return errors.Wrap(err, "could not decode result item")
			}
			callback(&result)
		}

		//This is checked during scan is added else if no item found break out of loop.
		if checkProgress {
			if items.Finished && len(items.Items) == 0 {
				break
			}
		} else if len(items.Items) == 0 {
			break
		}

		time.Sleep(pollInterval)
	}
	return nil
}

func (c *Client) GetScans() ([]GetScanRequest, error) {
	var items []GetScanRequest
	httpReq, err := retryablehttp.NewRequest(http.MethodGet, fmt.Sprintf("%s/scan", c.baseURL), nil)
	if err != nil {
		return items, errors.Wrap(err, "could not make request")
	}
	httpReq.Header.Set("X-API-Key", c.apiKey)

	resp, err := c.httpclient.Do(httpReq)
	if err != nil {
		return items, errors.Wrap(err, "could not make request.")
	}
	if err != nil {
		return items, errors.Wrap(err, "could not do get response.")
	}
	if resp.StatusCode != 200 {
		data, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return items, errors.Errorf("could not do request %d: %s", resp.StatusCode, string(data))
	}
	if err := jsoniter.NewDecoder(resp.Body).Decode(&items); err != nil {
		resp.Body.Close()
		return items, errors.Wrap(err, "could not decode results")
	}
	resp.Body.Close()

	return items, nil
}

//Delete a scan and it's issues by the scan id.
func (c *Client) DeleteScan(id string) (DeleteScanResults, error) {
	deletescan := DeleteScanResults{}
	httpReq, err := retryablehttp.NewRequest(http.MethodDelete, fmt.Sprintf("%s/scan?id=%s", c.baseURL, id), nil)
	if err != nil {
		return deletescan, errors.Wrap(err, "could not make request")
	}
	httpReq.Header.Set("X-API-Key", c.apiKey)

	resp, err := c.httpclient.Do(httpReq)
	if err != nil {
		return deletescan, errors.Wrap(err, "could not make request")
	}
	if err != nil {
		return deletescan, errors.Wrap(err, "could not do get result request")
	}
	if resp.StatusCode != 200 {
		data, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return deletescan, errors.Errorf("could not do request %d: %s", resp.StatusCode, string(data))
	}
	if err := jsoniter.NewDecoder(resp.Body).Decode(&deletescan); err != nil {
		resp.Body.Close()
		return deletescan, errors.Wrap(err, "could not delete scan")
	}
	resp.Body.Close()

	return deletescan, nil
}