fix(tencent):DescribeInstances support paging, fixes #1

404tk · 404tk · commit e6123f2b7287 · 2023-09-25T14:43:32.000+08:00
diff --git a/README.md b/README.md
@@ -1,6 +1,9 @@
 # cloudtoolkit
 Cloud Penetration Testing Toolkit
 
+## Usage  
+Reference [Wiki](https://github.com/404tk/cloudtoolkit/wiki)  
+
 ## Capability overview
 
 |          Providers          |                   Payload                   |                          Supported                           |
diff --git a/pkg/providers/tencent/cvm/instances.go b/pkg/providers/tencent/cvm/instances.go
@@ -27,7 +27,6 @@ func (d *Driver) NewClient() (*cvm.Client, error) {
 	return cvm.NewClient(d.Credential, region, cpf)
 }
 
-// GetResource returns all the resources in the store for a provider.
 func (d *Driver) GetResource(ctx context.Context) ([]schema.Host, error) {
 	list := schema.NewResources().Hosts
 	logger.Info("Start enumerating CVM ...")
@@ -48,45 +47,57 @@ func (d *Driver) GetResource(ctx context.Context) ([]schema.Host, error) {
 	}
 	flag := false
 	prevLength := 0
+	count := 0
 	for _, r := range regions {
 		d.Region = r
 		client, _ := d.NewClient()
 		request := cvm.NewDescribeInstancesRequest()
-		response, err := client.DescribeInstances(request)
-		if err != nil {
-			logger.Error("Enumerate CVM failed.")
-			return list, err
-		}
-
-		for _, instance := range response.Response.InstanceSet {
-			var ipv4, privateIPv4 string
-			if len(instance.PublicIpAddresses) > 0 {
-				ipv4 = *instance.PublicIpAddresses[0]
-			}
-			if len(instance.PrivateIpAddresses) > 0 {
-				privateIPv4 = *instance.PrivateIpAddresses[0]
+		request.Limit = common.Int64Ptr(100)
+		var offset int64 = 0
+		for {
+			request.Offset = common.Int64Ptr(offset)
+			response, err := client.DescribeInstances(request)
+			if err != nil {
+				logger.Error("Enumerate CVM failed.")
+				return list, err
 			}
-			host := schema.Host{
-				HostName:    *instance.InstanceName,
-				ID:          *instance.InstanceId,
-				PublicIPv4:  ipv4,
-				PrivateIpv4: privateIPv4,
-				Public:      ipv4 != "",
-				Region:      r,
+
+			for _, instance := range response.Response.InstanceSet {
+				var ipv4, privateIPv4 string
+				if len(instance.PublicIpAddresses) > 0 {
+					ipv4 = *instance.PublicIpAddresses[0]
+				}
+				if len(instance.PrivateIpAddresses) > 0 {
+					privateIPv4 = *instance.PrivateIpAddresses[0]
+				}
+				host := schema.Host{
+					HostName:    *instance.InstanceName,
+					ID:          *instance.InstanceId,
+					PublicIPv4:  ipv4,
+					PrivateIpv4: privateIPv4,
+					Public:      ipv4 != "",
+					Region:      r,
+				}
+				os_name := strings.Split(*instance.OsName, " ")[0]
+				if os_name == "Windows" {
+					host.OSType = "WINDOWS"
+				} else {
+					host.OSType = "LINUX_UNIX"
+				}
+				list = append(list, host)
 			}
-			os_name := strings.Split(*instance.OsName, " ")[0]
-			if os_name == "Windows" {
-				host.OSType = "WINDOWS"
-			} else {
-				host.OSType = "LINUX_UNIX"
+
+			if len(response.Response.InstanceSet) < 100 {
+				break
 			}
-			list = append(list, host)
+			offset += 100
 		}
 		select {
 		case <-ctx.Done():
 			goto done
 		default:
-			prevLength, flag = processbar.RegionPrint(r, len(response.Response.InstanceSet), prevLength, flag)
+			prevLength, flag = processbar.RegionPrint(r, len(list)-count, prevLength, flag)
+			count = len(list)
 		}
 	}
 done:
diff --git a/pkg/providers/tencent/lighthouse/instances.go b/pkg/providers/tencent/lighthouse/instances.go
@@ -26,7 +26,6 @@ func (d *Driver) NewClient() (*lighthouse.Client, error) {
 	return lighthouse.NewClient(d.Credential, region, cpf)
 }
 
-// GetResource returns all the resources in the store for a provider.
 func (d *Driver) GetResource(ctx context.Context) ([]schema.Host, error) {
 	list := schema.NewResources().Hosts
 	select {
@@ -53,40 +52,52 @@ func (d *Driver) GetResource(ctx context.Context) ([]schema.Host, error) {
 
 	flag := false
 	prevLength := 0
+	count := 0
 	for _, r := range regions {
 		d.Region = r
 		client, _ := d.NewClient()
 		request := lighthouse.NewDescribeInstancesRequest()
-		response, err := client.DescribeInstances(request)
-		if err != nil {
-			logger.Error("Enumerate Lighthouse failed.")
-			return list, err
-		}
-
-		for _, instance := range response.Response.InstanceSet {
-			var ipv4, privateIPv4 string
-			if len(instance.PublicAddresses) > 0 {
-				ipv4 = *instance.PublicAddresses[0]
+		request.Limit = common.Int64Ptr(100)
+		var offset int64 = 0
+		for {
+			request.Offset = common.Int64Ptr(offset)
+			response, err := client.DescribeInstances(request)
+			if err != nil {
+				logger.Error("Enumerate Lighthouse failed.")
+				return list, err
 			}
-			if len(instance.PrivateAddresses) > 0 {
-				privateIPv4 = *instance.PrivateAddresses[0]
+
+			for _, instance := range response.Response.InstanceSet {
+				var ipv4, privateIPv4 string
+				if len(instance.PublicAddresses) > 0 {
+					ipv4 = *instance.PublicAddresses[0]
+				}
+				if len(instance.PrivateAddresses) > 0 {
+					privateIPv4 = *instance.PrivateAddresses[0]
+				}
+				_host := schema.Host{
+					HostName:    *instance.InstanceName,
+					ID:          *instance.InstanceId,
+					PublicIPv4:  ipv4,
+					PrivateIpv4: privateIPv4,
+					OSType:      *instance.PlatformType, // LINUX_UNIX or WINDOWS
+					Public:      ipv4 != "",
+					Region:      r,
+				}
+				list = append(list, _host)
 			}
-			_host := schema.Host{
-				HostName:    *instance.InstanceName,
-				ID:          *instance.InstanceId,
-				PublicIPv4:  ipv4,
-				PrivateIpv4: privateIPv4,
-				OSType:      *instance.PlatformType, // LINUX_UNIX or WINDOWS
-				Public:      ipv4 != "",
-				Region:      r,
+
+			if len(response.Response.InstanceSet) < 100 {
+				break
 			}
-			list = append(list, _host)
+			offset += 100
 		}
 		select {
 		case <-ctx.Done():
 			goto done
 		default:
-			prevLength, flag = processbar.RegionPrint(r, len(response.Response.InstanceSet), prevLength, flag)
+			prevLength, flag = processbar.RegionPrint(r, len(list)-count, prevLength, flag)
+			count = len(list)
 		}
 	}
 done:
