diff --git a/device.go b/device.go
index 4816d7e..13f44c9 100644
--- a/device.go
+++ b/device.go
@@ -657,7 +657,7 @@ func (d *device) Reboot() (err error) {
 	return
 }
 
-func (d *device) PowerSource() (powerInfo interface{}, err error) {
+func (d *device) PowerSource() (powerInfo map[string]interface{}, err error) {
 	if _, err = d.lockdownService(); err != nil {
 		return
 	}
diff --git a/diagnosticsrelay.go b/diagnosticsrelay.go
index 0c80f0d..6cb75ff 100644
--- a/diagnosticsrelay.go
+++ b/diagnosticsrelay.go
@@ -17,7 +17,7 @@ type diagnostics struct {
 func (d *diagnostics) Reboot() (err error) {
 	var pkt libimobiledevice.Packet
 	if pkt, err = d.client.NewXmlPacket(
-		d.client.NewBasicRequest("Restart"),
+		d.client.NewBasicRequest("Restart", nil),
 	); err != nil {
 		return
 	}
@@ -30,7 +30,7 @@ func (d *diagnostics) Reboot() (err error) {
 func (d *diagnostics) Shutdown() (err error) {
 	var pkt libimobiledevice.Packet
 	if pkt, err = d.client.NewXmlPacket(
-		d.client.NewBasicRequest("Shutdown"),
+		d.client.NewBasicRequest("Shutdown", nil),
 	); err != nil {
 		return
 	}
@@ -40,14 +40,11 @@ func (d *diagnostics) Shutdown() (err error) {
 	return
 }
 
-func (d *diagnostics) PowerSource() (powerInfo interface{}, err error) {
-	var data map[string]string
-	data["Request"] = "IORegistry"
-	data["EntryClass"] = "IOPMPowerSource"
-	data["Label"] = "org.cloud.sonic.gidevice"
+func (d *diagnostics) PowerSource() (powerInfo map[string]interface{}, err error) {
 	var pkt libimobiledevice.Packet
+	ioRegistry := "IOPMPowerSource"
 	if pkt, err = d.client.NewXmlPacket(
-		d.client.NewBasicRequest("Shutdown"),
+		d.client.NewBasicRequest("IORegistry", &ioRegistry),
 	); err != nil {
 		return nil, err
 	}
@@ -57,6 +54,9 @@ func (d *diagnostics) PowerSource() (powerInfo interface{}, err error) {
 	if pkt, err = d.client.ReceivePacket(); err != nil {
 		return nil, err
 	}
+	if powerInfo == nil {
+		powerInfo = make(map[string]interface{})
+	}
 	if err = pkt.Unmarshal(powerInfo); err != nil {
 		return nil, err
 	}
diff --git a/idevice.go b/idevice.go
index 7f07dcf..f5cac6a 100644
--- a/idevice.go
+++ b/idevice.go
@@ -67,7 +67,7 @@ type Device interface {
 
 	Reboot() error
 	Shutdown() error
-	PowerSource() (powerInfo interface{}, err error)
+	PowerSource() (powerInfo map[string]interface{}, err error)
 
 	crashReportMoverService() (crashReportMover CrashReportMover, err error)
 	MoveCrashReport(hostDir string, opts ...CrashReportMoverOption) (err error)
@@ -235,7 +235,7 @@ type Perfd interface {
 type DiagnosticsRelay interface {
 	Reboot() error
 	Shutdown() error
-	PowerSource() (powerInfo interface{}, err error)
+	PowerSource() (powerInfo map[string]interface{}, err error)
 }
 
 type WebInspector interface {
diff --git a/pkg/libimobiledevice/diagnosticsrelay.go b/pkg/libimobiledevice/diagnosticsrelay.go
index 0ff0f3b..38e487d 100644
--- a/pkg/libimobiledevice/diagnosticsrelay.go
+++ b/pkg/libimobiledevice/diagnosticsrelay.go
@@ -5,8 +5,9 @@ const (
 )
 
 type DiagnosticsRelayBasicRequest struct {
-	Request string `plist:"Request"`
-	Label   string `plist:"Label"`
+	Request    string  `plist:"Request"`
+	Label      string  `plist:"Label"`
+	EntryClass *string `plist:"EntryClass,omitempty"`
 }
 
 func NewDiagnosticsRelayClient(innerConn InnerConn) *DiagnosticsRelayClient {
@@ -23,10 +24,11 @@ func (c *DiagnosticsRelayClient) InnerConn() InnerConn {
 	return c.client.innerConn
 }
 
-func (c *DiagnosticsRelayClient) NewBasicRequest(relayType string) *DiagnosticsRelayBasicRequest {
+func (c *DiagnosticsRelayClient) NewBasicRequest(relayType string, entryClass *string) *DiagnosticsRelayBasicRequest {
 	return &DiagnosticsRelayBasicRequest{
-		Request: relayType,
-		Label:   BundleID,
+		Request:    relayType,
+		Label:      BundleID,
+		EntryClass: entryClass,
 	}
 }