From 1c10978facb5f2dc2ede98392714ce57f6ccc19a Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Fri, 14 Apr 2023 23:50:40 +0800
Subject: [PATCH] fix: screencap compatibility for shell v1 and v2

---
 docs/CHANGELOG.md            |  1 +
 hrp/internal/version/VERSION |  2 +-
 hrp/pkg/gadb/device.go       | 17 ++++++++++++++++-
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index d7b5f38d..35632507 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -13,6 +13,7 @@
 - fix: stop logcat only when enabled
 - fix: do not fail case when kill logcat error
 - fix: take screenshot after each step
+- fix: screencap compatibility for shell v1 and v2
 
 ## v4.3.2 (2022-12-26)
 
diff --git a/hrp/internal/version/VERSION b/hrp/internal/version/VERSION
index 45af2229..f907911d 100644
--- a/hrp/internal/version/VERSION
+++ b/hrp/internal/version/VERSION
@@ -1 +1 @@
-v4.3.3.2304142149
\ No newline at end of file
+v4.3.3.2304142356
\ No newline at end of file
diff --git a/hrp/pkg/gadb/device.go b/hrp/pkg/gadb/device.go
index 87a496ea..181609b0 100644
--- a/hrp/pkg/gadb/device.go
+++ b/hrp/pkg/gadb/device.go
@@ -572,5 +572,20 @@ func (d *Device) Uninstall(packageName string, keepData ...bool) (string, error)
 }
 
 func (d *Device) ScreenCap() ([]byte, error) {
-	return d.RunShellCommandWithBytes("screencap", "-p")
+	if d.HasFeature(FeatShellV2) {
+		return d.RunShellCommandV2WithBytes("screencap", "-p")
+	}
+
+	// for shell v1, screenshot buffer maybe truncated
+	// thus we firstly save it to local file and then pull it
+	tempPath := fmt.Sprintf("/data/local/tmp/screenshot_%d.png",
+		time.Now().Unix())
+	_, err := d.RunShellCommandWithBytes("screencap", "-p", tempPath)
+	if err != nil {
+		return nil, err
+	}
+
+	buffer := bytes.NewBuffer(nil)
+	err = d.Pull(tempPath, buffer)
+	return buffer.Bytes(), err
 }