god-eye/internal/agent/tools.go

package agent

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"time"

	godns "god-eye/internal/dns"
)

// --- built-in tools -------------------------------------------------------
//
// These tools cover the minimum needed for a planner to investigate
// discovered hosts without reinventing basic primitives. Fase 3 workers
// receive curated subsets via Toolset.

// HTTPRequestTool fetches an arbitrary URL and returns status, headers,
// and (truncated) body. Maximum 64KB body returned.
type HTTPRequestTool struct {
	Client *http.Client
}

func NewHTTPRequestTool(timeoutSec int) *HTTPRequestTool {
	return &HTTPRequestTool{
		Client: &http.Client{
			Timeout: time.Duration(timeoutSec) * time.Second,
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		},
	}
}

func (t *HTTPRequestTool) Name() string        { return "http_request" }
func (t *HTTPRequestTool) Description() string { return "Fetch an HTTP(S) URL and return status + headers + first 64KB of body." }

func (t *HTTPRequestTool) Schema() map[string]interface{} {
	return map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"url":    map[string]interface{}{"type": "string"},
			"method": map[string]interface{}{"type": "string", "default": "GET"},
			"headers": map[string]interface{}{
				"type":                 "object",
				"additionalProperties": map[string]interface{}{"type": "string"},
			},
		},
		"required": []string{"url"},
	}
}

func (t *HTTPRequestTool) Call(ctx context.Context, args map[string]interface{}) (string, error) {
	url, _ := args["url"].(string)
	if url == "" {
		return "", errors.New("url is required")
	}
	method, _ := args["method"].(string)
	if method == "" {
		method = "GET"
	}
	req, err := http.NewRequestWithContext(ctx, method, url, nil)
	if err != nil {
		return "", err
	}
	if hdrs, ok := args["headers"].(map[string]interface{}); ok {
		for k, v := range hdrs {
			if s, ok := v.(string); ok {
				req.Header.Set(k, s)
			}
		}
	}
	req.Header.Set("User-Agent", "god-eye-v2-agent")

	resp, err := t.Client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
	out := map[string]interface{}{
		"status_code": resp.StatusCode,
		"headers":     flattenHeaders(resp.Header),
		"body":        string(body),
	}
	b, _ := json.Marshal(out)
	return string(b), nil
}

// DNSResolveTool resolves a hostname to A/CNAME/PTR records.
type DNSResolveTool struct {
	Resolvers []string
	TimeoutSec int
}

func NewDNSResolveTool(resolvers []string, timeoutSec int) *DNSResolveTool {
	if len(resolvers) == 0 {
		resolvers = []string{"8.8.8.8:53", "1.1.1.1:53"}
	}
	return &DNSResolveTool{Resolvers: resolvers, TimeoutSec: timeoutSec}
}

func (t *DNSResolveTool) Name() string        { return "dns_resolve" }
func (t *DNSResolveTool) Description() string { return "Resolve a hostname to A/CNAME/PTR records." }

func (t *DNSResolveTool) Schema() map[string]interface{} {
	return map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"hostname": map[string]interface{}{"type": "string"},
		},
		"required": []string{"hostname"},
	}
}

func (t *DNSResolveTool) Call(_ context.Context, args map[string]interface{}) (string, error) {
	name, _ := args["hostname"].(string)
	if name == "" {
		return "", errors.New("hostname is required")
	}
	ips := godns.ResolveSubdomain(name, t.Resolvers, t.TimeoutSec)
	cname := godns.ResolveCNAME(name, t.Resolvers, t.TimeoutSec)
	out := map[string]interface{}{"ips": ips, "cname": cname}
	b, _ := json.Marshal(out)
	return string(b), nil
}

func flattenHeaders(h http.Header) map[string]string {
	out := make(map[string]string, len(h))
	for k, vs := range h {
		if len(vs) == 0 {
			continue
		}
		out[k] = vs[0]
	}
	return out
}