From f624ef4dcac0638177e822ab3886a293873ccca9 Mon Sep 17 00:00:00 2001
From: blacktop <blacktop@users.noreply.github.com>
Date: Wed, 18 Sep 2024 12:30:35 -0600
Subject: [PATCH] fix: `ipsw dtree` cmd to work on local IPSWs as input arg

---
 cmd/ipsw/cmd/device_tree.go | 42 ++++++++++++++++++++++++++-----------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/cmd/ipsw/cmd/device_tree.go b/cmd/ipsw/cmd/device_tree.go
index 369d44c26..04a4fdda3 100644
--- a/cmd/ipsw/cmd/device_tree.go
+++ b/cmd/ipsw/cmd/device_tree.go
@@ -22,10 +22,12 @@ THE SOFTWARE.
 package cmd
 
 import (
+	"archive/zip"
 	"bytes"
 	"encoding/json"
 	"fmt"
 	"os"
+	"path/filepath"
 
 	// "sort"
 
@@ -34,7 +36,6 @@ import (
 	"github.com/blacktop/ipsw/internal/magic"
 	"github.com/blacktop/ipsw/internal/utils"
 	"github.com/blacktop/ipsw/pkg/devicetree"
-	"github.com/pkg/errors"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -82,34 +83,51 @@ var deviceTreeCmd = &cobra.Command{
 			}
 			dtrees, err = devicetree.ParseZipFiles(zr.File)
 			if err != nil {
-				return errors.Wrap(err, "failed to extract DeviceTree")
+				return fmt.Errorf("failed to extract DeviceTree: %v", err)
 			}
 		} else {
 			var dtree *devicetree.DeviceTree
 
-			content, err := os.ReadFile(args[0])
-			if err != nil {
-				return errors.Wrap(err, "failed to read DeviceTree")
-			}
-
-			if ok, _ := magic.IsImg3(args[0]); ok {
+			if ok, _ := magic.IsZip(filepath.Clean(args[0])); ok {
+				zr, err := zip.OpenReader(args[0])
+				if err != nil {
+					return fmt.Errorf("failed to open zip: %v", err)
+				}
+				dtrees, err = devicetree.ParseZipFiles(zr.File)
+				if err != nil {
+					return fmt.Errorf("failed to extract DeviceTree: %v", err)
+				}
+			} else if ok, _ := magic.IsImg3(args[0]); ok {
+				content, err := os.ReadFile(args[0])
+				if err != nil {
+					return fmt.Errorf("failed to read DeviceTree: %v", err)
+				}
 				dtree, err = devicetree.ParseImg3Data(content)
 				if err != nil {
-					return errors.Wrap(err, "failed to extract DeviceTree")
+					return fmt.Errorf("failed to extract DeviceTree: %v", err)
 				}
+				dtrees[args[0]] = dtree
 			} else if ok, _ := magic.IsIm4p(args[0]); ok {
+				content, err := os.ReadFile(args[0])
+				if err != nil {
+					return fmt.Errorf("failed to read DeviceTree: %v", err)
+				}
 				dtree, err = devicetree.ParseImg4Data(content)
 				if err != nil {
-					return errors.Wrap(err, "failed to extract DeviceTree")
+					return fmt.Errorf("failed to extract DeviceTree: %v", err)
 				}
+				dtrees[args[0]] = dtree
 			} else {
+				content, err := os.ReadFile(args[0])
+				if err != nil {
+					return fmt.Errorf("failed to read DeviceTree: %v", err)
+				}
 				dtree, err = devicetree.ParseData(bytes.NewReader(content))
 				if err != nil {
 					return fmt.Errorf("failed to parse DeviceTree: %v", err)
 				}
+				dtrees[args[0]] = dtree
 			}
-
-			dtrees[args[0]] = dtree
 		}
 
 		for name, dtree := range dtrees {