heicode/heicode/controller/aws_resource_discovery_services.go

package controller

import (
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/heicode/manager/common"
)

// Service-level AWS discovery (#62): augments the tagged-only GetResources pass
// (aws_resource_discovery.go) with direct List/Describe so **untagged** EC2/RDS/S3/Lambda
// resources are found too — closing the "tagged only ≠ full account" gap (#5 follow-up).
//
// 设计:
//   - discoverAWS 编排:先 tagged(最全的标签信息),再各服务级 List/Describe,按 ARN 合并去重
//     (tagged 优先保留,信息更丰富)。
//   - **失败降级**:任一服务调用失败(如缺该服务读权限 / AccessDenied)只记日志并跳过,绝不让
//     整次发现失败;仅当「tagged 报错且无任何服务返回结果」时,才把 tagged 的原始错误抛出,
//     好让账号凭据/区域问题对用户可见。
//   - 复用 awsSigV4Authorization(已对 AWS get-vanilla 向量验证)签 GET;无 AWS SDK 依赖。
//   - 区域:EC2/RDS/Lambda 用账号配置的 region;S3 ListBuckets 是全局接口(以 us-east-1 签名)。
//     跨 region 全量扫描不在本次范围(EC2/RDS 仅扫配置 region),由 #62 备注另行跟踪。
//   - 解析逻辑(parseAWS* )为纯函数,便于用真实响应样本单测。

// discoverAWS enumerates the account's resources by merging the tagged-resources pass
// with service-level List/Describe (#62). Replaces the old tagged-only discoverAWS.
func discoverAWS(creds map[string]any, now time.Time, httpClient *http.Client) ([]discoveredCloudResource, error) {
	cred := awsDiscoveryCredentials{
		AccessKeyID:     mapString(creds, "access_key_id"),
		SecretAccessKey: mapString(creds, "secret_access_key"),
		SessionToken:    mapString(creds, "session_token"),
		Region:          mapString(creds, "region"),
	}
	if err := cred.validate(); err != nil {
		return nil, err
	}

	var lists [][]discoveredCloudResource
	tagged, taggedErr := discoverAWSTagged(cred, now, httpClient)
	if taggedErr != nil {
		common.SysLog("discoverAWS tagged: " + taggedErr.Error())
	} else {
		lists = append(lists, tagged) // tagged 先入,合并时优先保留
	}

	services := []struct {
		name string
		fn   func(awsDiscoveryCredentials, time.Time, *http.Client) ([]discoveredCloudResource, error)
	}{
		{"ec2", discoverAWSEC2},
		{"rds", discoverAWSRDS},
		{"s3", discoverAWSS3},
		{"lambda", discoverAWSLambda},
	}
	for _, svc := range services {
		list, err := svc.fn(cred, now, httpClient)
		if err != nil {
			common.SysLog("discoverAWS " + svc.name + ": " + err.Error())
			continue
		}
		lists = append(lists, list)
	}

	merged := mergeDedupeByExternalID(lists...)
	// 全部为空且 tagged 当初报错 → 抛出原始错误,让凭据/区域问题可见(纯权限缺失则降级返回已得结果)。
	if len(merged) == 0 && taggedErr != nil {
		return nil, taggedErr
	}
	return merged, nil
}

// mergeDedupeByExternalID 按 ExternalId(ARN)合并多组发现结果,先到先得(靠前的列表优先保留),
// 保持稳定顺序。纯函数,便于单测。
func mergeDedupeByExternalID(lists ...[]discoveredCloudResource) []discoveredCloudResource {
	seen := map[string]bool{}
	out := make([]discoveredCloudResource, 0)
	for _, list := range lists {
		for _, d := range list {
			id := strings.TrimSpace(d.ExternalId)
			if id == "" || seen[id] {
				continue
			}
			seen[id] = true
			out = append(out, d)
		}
	}
	return out
}

// awsSignedGet 发起一次 SigV4 签名的 GET(空 body),返回响应体。非 2xx 视为错误。
func awsSignedGet(cred awsDiscoveryCredentials, rawURL, service string, now time.Time, httpClient *http.Client) ([]byte, error) {
	auth, amzDate, err := awsSigV4Authorization(http.MethodGet, rawURL, map[string]string{}, nil, cred, service, now)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest(http.MethodGet, rawURL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("X-Amz-Date", amzDate)
	req.Header.Set("Authorization", auth)
	if strings.TrimSpace(cred.SessionToken) != "" {
		req.Header.Set("X-Amz-Security-Token", cred.SessionToken)
	}
	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 16<<20))
	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
		return body, fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	return body, nil
}

// ---- EC2 DescribeInstances (Query 协议, XML) ----

type ec2DescribeInstancesResponse struct {
	XMLName      xml.Name `xml:"DescribeInstancesResponse"`
	Reservations []struct {
		OwnerID   string `xml:"ownerId"`
		Instances []struct {
			InstanceID   string `xml:"instanceId"`
			InstanceType string `xml:"instanceType"`
			State        struct {
				Name string `xml:"name"`
			} `xml:"instanceState"`
			Placement struct {
				AvailabilityZone string `xml:"availabilityZone"`
			} `xml:"placement"`
			Tags []struct {
				Key   string `xml:"key"`
				Value string `xml:"value"`
			} `xml:"tagSet>item"`
		} `xml:"instancesSet>item"`
	} `xml:"reservationSet>item"`
	NextToken string `xml:"nextToken"`
}

// parseEC2Instances 解析 DescribeInstances XML → 归一化资源 + nextToken。纯函数。
func parseEC2Instances(body []byte, region string) ([]discoveredCloudResource, string, error) {
	var r ec2DescribeInstancesResponse
	if err := xml.Unmarshal(body, &r); err != nil {
		return nil, "", err
	}
	out := []discoveredCloudResource{}
	for _, res := range r.Reservations {
		for _, inst := range res.Instances {
			if strings.TrimSpace(inst.InstanceID) == "" {
				continue
			}
			arn := fmt.Sprintf("arn:aws:ec2:%s:%s:instance/%s", region, res.OwnerID, inst.InstanceID)
			tags := map[string]any{}
			name := inst.InstanceID
			for _, t := range inst.Tags {
				tags[t.Key] = t.Value
				if strings.EqualFold(t.Key, "Name") && strings.TrimSpace(t.Value) != "" {
					name = t.Value
				}
			}
			out = append(out, discoveredCloudResource{
				ExternalId:   arn,
				Name:         name,
				NativeType:   "AWS::EC2::Instance",
				Location:     region,
				BindingScope: "aws:" + arn,
				Metadata: map[string]any{
					"region":        region,
					"instance_type": inst.InstanceType,
					"state":         inst.State.Name,
					"tags":          tags,
				},
			})
		}
	}
	return out, strings.TrimSpace(r.NextToken), nil
}

func discoverAWSEC2(cred awsDiscoveryCredentials, now time.Time, httpClient *http.Client) ([]discoveredCloudResource, error) {
	out := []discoveredCloudResource{}
	nextToken := ""
	for page := 0; page < 50; page++ {
		q := url.Values{}
		q.Set("Action", "DescribeInstances")
		q.Set("Version", "2016-11-15")
		q.Set("MaxResults", "100")
		if nextToken != "" {
			q.Set("NextToken", nextToken)
		}
		rawURL := fmt.Sprintf("https://ec2.%s.amazonaws.com/?%s", cred.Region, q.Encode())
		body, err := awsSignedGet(cred, rawURL, "ec2", now, httpClient)
		if err != nil {
			return nil, err
		}
		page1, token, err := parseEC2Instances(body, cred.Region)
		if err != nil {
			return nil, err
		}
		out = append(out, page1...)
		nextToken = token
		if nextToken == "" {
			break
		}
	}
	return out, nil
}

// ---- RDS DescribeDBInstances (Query 协议, XML) ----

type rdsDescribeDBInstancesResponse struct {
	XMLName xml.Name `xml:"DescribeDBInstancesResponse"`
	Result  struct {
		Marker      string `xml:"Marker"`
		DBInstances []struct {
			Arn              string `xml:"DBInstanceArn"`
			Identifier       string `xml:"DBInstanceIdentifier"`
			Engine           string `xml:"Engine"`
			Status           string `xml:"DBInstanceStatus"`
			AvailabilityZone string `xml:"AvailabilityZone"`
		} `xml:"DBInstances>DBInstance"`
	} `xml:"DescribeDBInstancesResult"`
}

// parseRDSInstances 解析 DescribeDBInstances XML → 归一化资源 + marker。纯函数。
func parseRDSInstances(body []byte, region string) ([]discoveredCloudResource, string, error) {
	var r rdsDescribeDBInstancesResponse
	if err := xml.Unmarshal(body, &r); err != nil {
		return nil, "", err
	}
	out := []discoveredCloudResource{}
	for _, db := range r.Result.DBInstances {
		arn := strings.TrimSpace(db.Arn)
		if arn == "" {
			continue // 无 ARN 不入(稳定外部 id 缺失)
		}
		name := strings.TrimSpace(db.Identifier)
		if name == "" {
			name = arn
		}
		out = append(out, discoveredCloudResource{
			ExternalId:   arn,
			Name:         name,
			NativeType:   "AWS::RDS::DBInstance",
			Location:     region,
			BindingScope: "aws:" + arn,
			Metadata: map[string]any{
				"region": region,
				"engine": db.Engine,
				"status": db.Status,
				"az":     db.AvailabilityZone,
			},
		})
	}
	return out, strings.TrimSpace(r.Result.Marker), nil
}

func discoverAWSRDS(cred awsDiscoveryCredentials, now time.Time, httpClient *http.Client) ([]discoveredCloudResource, error) {
	out := []discoveredCloudResource{}
	marker := ""
	for page := 0; page < 50; page++ {
		q := url.Values{}
		q.Set("Action", "DescribeDBInstances")
		q.Set("Version", "2014-10-31")
		q.Set("MaxRecords", "100")
		if marker != "" {
			q.Set("Marker", marker)
		}
		rawURL := fmt.Sprintf("https://rds.%s.amazonaws.com/?%s", cred.Region, q.Encode())
		body, err := awsSignedGet(cred, rawURL, "rds", now, httpClient)
		if err != nil {
			return nil, err
		}
		page1, next, err := parseRDSInstances(body, cred.Region)
		if err != nil {
			return nil, err
		}
		out = append(out, page1...)
		marker = next
		if marker == "" {
			break
		}
	}
	return out, nil
}

// ---- S3 ListBuckets (REST-XML, 全局) ----

type s3ListAllMyBucketsResult struct {
	XMLName xml.Name `xml:"ListAllMyBucketsResult"`
	Buckets []struct {
		Name string `xml:"Name"`
	} `xml:"Buckets>Bucket"`
}

// parseS3Buckets 解析 ListBuckets XML → 归一化资源。纯函数。S3 bucket ARN 无 region/account 段。
func parseS3Buckets(body []byte) ([]discoveredCloudResource, error) {
	var r s3ListAllMyBucketsResult
	if err := xml.Unmarshal(body, &r); err != nil {
		return nil, err
	}
	out := []discoveredCloudResource{}
	for _, b := range r.Buckets {
		name := strings.TrimSpace(b.Name)
		if name == "" {
			continue
		}
		arn := "arn:aws:s3:::" + name
		out = append(out, discoveredCloudResource{
			ExternalId:   arn,
			Name:         name,
			NativeType:   "AWS::S3::Bucket",
			Location:     "", // 全局列举不含 region;按需 GetBucketLocation(未做,避免逐桶请求)
			BindingScope: "aws:" + arn,
			Metadata:     map[string]any{"bucket": name},
		})
	}
	return out, nil
}

func discoverAWSS3(cred awsDiscoveryCredentials, now time.Time, httpClient *http.Client) ([]discoveredCloudResource, error) {
	// ListBuckets 是全局接口,固定以 us-east-1 / service=s3 签名。
	s3Cred := cred
	s3Cred.Region = "us-east-1"
	body, err := awsSignedGet(s3Cred, "https://s3.amazonaws.com/", "s3", now, httpClient)
	if err != nil {
		return nil, err
	}
	return parseS3Buckets(body)
}

// ---- Lambda ListFunctions (REST-JSON) ----

type lambdaListFunctionsResponse struct {
	Functions []struct {
		FunctionArn  string `json:"FunctionArn"`
		FunctionName string `json:"FunctionName"`
		Runtime      string `json:"Runtime"`
	} `json:"Functions"`
	NextMarker string `json:"NextMarker"`
}

// parseLambdaFunctions 解析 ListFunctions JSON → 归一化资源 + nextMarker。纯函数。
func parseLambdaFunctions(body []byte, region string) ([]discoveredCloudResource, string, error) {
	var r lambdaListFunctionsResponse
	if err := common.Unmarshal(body, &r); err != nil {
		return nil, "", err
	}
	out := []discoveredCloudResource{}
	for _, fn := range r.Functions {
		arn := strings.TrimSpace(fn.FunctionArn)
		if arn == "" {
			continue
		}
		name := strings.TrimSpace(fn.FunctionName)
		if name == "" {
			name = arn
		}
		out = append(out, discoveredCloudResource{
			ExternalId:   arn,
			Name:         name,
			NativeType:   "AWS::Lambda::Function",
			Location:     region,
			BindingScope: "aws:" + arn,
			Metadata:     map[string]any{"region": region, "runtime": fn.Runtime},
		})
	}
	return out, strings.TrimSpace(r.NextMarker), nil
}

func discoverAWSLambda(cred awsDiscoveryCredentials, now time.Time, httpClient *http.Client) ([]discoveredCloudResource, error) {
	out := []discoveredCloudResource{}
	marker := ""
	for page := 0; page < 50; page++ {
		q := url.Values{}
		q.Set("MaxItems", "50")
		if marker != "" {
			q.Set("Marker", marker)
		}
		rawURL := fmt.Sprintf("https://lambda.%s.amazonaws.com/2015-03-31/functions/?%s", cred.Region, q.Encode())
		body, err := awsSignedGet(cred, rawURL, "lambda", now, httpClient)
		if err != nil {
			return nil, err
		}
		page1, next, err := parseLambdaFunctions(body, cred.Region)
		if err != nil {
			return nil, err
		}
		out = append(out, page1...)
		marker = next
		if marker == "" {
			break
		}
	}
	return out, nil
}