diagnostics.py

#!/usr/bin/python
# coding: utf-8

# Copyright (c) 2013 Mountainstorm
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


from MobileDevice import *
from amdevice import *
from plistservice import *


class Diagnostics(PlistService):
    u'''Communicates with the diagnostic service and allows retrival of
    diagnostic information
    '''

    TYPE_ALL = u'All'
    TYPE_WIFI = u'WiFi'
    TYPE_GASGAUGE = u'GasGauge'
    TYPE_NAND = u'NAND'

    ACTION_WAIT_FOR_DISCONNECT = u'WaitForDisconnect'
    ACTION_DISPLAY_PASS = u'DisplayPass'
    ACTION_DISPLAY_FAIL = u'DisplayFail'

    gestalt_names = [
    ]

    def __init__(self, amdevice):
        PlistService.__init__(
            self,
            amdevice,
            [AMSVC_MOBILE_DIAGNOSTICS_RELAY, AMSVC_IOSDIAGNOSTICS_RELAY],
            kCFPropertyListXMLFormat_v1_0
        )

    def disconnect(self):
        PlistService.disconnect(self)

    def _action(self, name, actions):
        msg = {
            u'Request': name
        }
        for action in actions:
            msg[action] = True
        self._sendmsg(msg)
        return self._chkresult(self._recvmsg())

    def _chkresult(self, result):
        retval = False
        if u'Status' in result and result[u'Status'] == u'Success':
            retval = True
        return retval

    def sleep(self):
        u'''Sends the device to sleep; disconnecting it from the host'''
        self._sendmsg({u'Request': u'Sleep'})
        return self._chkresult(self._recvmsg())

    def restart(self, actions=[]):
        u'''Restarts the device'''
        return self._action(u'Restart', actions)

    def shutdown(self, actions=[]):
        u'''Shutsdownt the host'''
        return self._action(u'Shutdown', actions)

    def diagnostics(self, typ):
        u'''Retrieves diagnostic information from the device; i.e. gasguage,
        wifi, nand.

        Arguments:
        typ -- one of the ACTION_* defines
        '''
        retval = None
        self._sendmsg({u'Request': typ})
        msg = self._recvmsg()
        if self._chkresult(msg):
            retval = msg[u'Diagnostics']
        return retval

    # look in com.apple.mobilegestalt.plist for values
    def mobilegestalt(self, keys):
        u'''Retrieves mobile gestalt values

        Arguments:
        keys -- an array of keys to retrieve
        '''
        self._sendmsg({
            u'Request': u'MobileGestalt',
            u'MobileGestaltKeys': keys
        })
        msg = self._recvmsg()
        if self._chkresult(msg):
            retval = msg[u'Diagnostics'][u'MobileGestalt']
        return retval

    # IODeviceTree, IOPower, IOService
    def ioregistry(self, plane=u'IOService', name=None, cls=None):
        u'''Retrieves IORegistry keys and values

        Arguments:
        plane -- a IOReg plane name; IODeviceTree, IOPower, IOService
        name -- the name of the key to retrieve
        cls -- the class of the key to retrieve
        '''
        msg = {
            u'Request': u'IORegistry',
            u'CurrentPlane': plane
        }
        if name is not None:
            msg[u'EntryName'] = name
        if cls is not None:
            msg[u'EntryClass'] = cls

        self._sendmsg(msg)
        msg = self._recvmsg()
        if self._chkresult(msg):
            retval = msg[u'Diagnostics']
        return retval


def register_argparse_diag(cmdargs):
    import argparse
    import pprint
    import sys

    def ioreg_dmp(d, pad=u'  '):
        # XXX should really adjust to get rid of all the -2's
        if isinstance(d, list):
            sys.stdout.write(u'[\n')
            for i in range(len(d)):
                sys.stdout.write(pad)
                ioreg_dmp(d[i], pad + u'  ')
            sys.stdout.write(pad[:-2] + u']\n')
        elif isinstance(d, dict):
            i = 0
            if u'regEntry' in d:
                # its a reg entry
                inheritance = d[u'inheritance'].replace(u' : ', u':')
                sys.stdout.write(
                    u'+-o ' + d[u'name'] + u'  <class ' + inheritance + u', ' + d[u'state'] + u'>\n'
                )
                op = pad
                if len(d[u'children']) > 0:
                    pad += u'| '
                else:
                    pad += u'  '
                sys.stdout.write(pad)
                ioreg_dmp(d[u'properties'], pad + u'  ')
                sys.stdout.write(pad + u'\n')
                i = 0
                for c in d[u'children']:
                    if i == len(d[u'children'])-1:
                        pad = pad[:-2] + u'  '
                    sys.stdout.write(pad[:-2])
                    ioreg_dmp(c, pad)
                    i += 1
            else:
                # something else?
                sys.stdout.write(u'{\n')
                for k, v in d.iteritems():
                    sys.stdout.write(pad + u'"' + k + u'" = ')
                    ioreg_dmp(v, pad + u'  ')
                sys.stdout.write(pad[:-2] + u'}\n')
        elif isinstance(d, unicode):
            sys.stdout.write(u'"' + d + u'"\n')
        elif isinstance(d, str):
            # dont know why but some strings name/inheritence etc come back as
            # CFData rather then CFString's - so we'll try to convert; if it
            # doesn't look like a string
            isdata = False
            nullterm = False
            for c in d:
                if nullterm == True and ord(c) != 0:
                    isdata = True
                    break
                if ord(c) == 0:
                    nullterm = True
                elif ord(c) < 32 or ord(c) > 127:
                    isdata = True
                    break
            if isdata:
                sys.stdout.write(d.encode(u'hex') + u'\n')
            else:
                sys.stdout.write(u'"' + d + u'"\n')
        elif isinstance(d, (int, long, float, complex)):
            sys.stdout.write(d.__str__() + u'\n')
        else:
            print type(d)
            pprint.pprint(d)

    def cmd_gestalt(args, dev):
        diag = Diagnostics(dev)
        if args.name is not None:
            print(diag.mobilegestalt([args.name.decode(u'utf-8')]))
        else:
            print(diag.mobilegestalt(Diagnostics.gestalt_names))
        diag.disconnect()

    def cmd_all(args, dev):
        diag = Diagnostics(dev)
        pprint.pprint(diag.diagnostics(Diagnostics.TYPE_ALL))
        diag.disconnect()

    def cmd_ioreg(args, dev):
        diag = Diagnostics(dev)
        plane = u'IOService'
        name = None
        cls = None
        # XXX check the name/cls params actually work
        if args.plane is not None:
            plane = args.plane.decode(u'utf-8')
        if args.name is not None:
            name = args.name.decode(u'utf-8')
        if args.cls is not None:
            cls = args.cls.decode(u'utf-8')
        ioreg_dmp(diag.ioregistry(plane, name, cls)[u'IORegistry'])
        diag.disconnect()

    def cmd_sleep(args, dev):
        diag = Diagnostics(dev)
        diag.sleep()
        diag.disconnect()

    def cmd_restart(args, dev):
        diag = Diagnostics(dev)
        if args.quick:
            diag.restart([Diagnostics.ACTION_DISPLAY_FAIL])
        else:
            diag.restart([Diagnostics.ACTION_DISPLAY_PASS])
        diag.disconnect()

    def cmd_shutdown(args, dev):
        diag = Diagnostics(dev)
        print args.quick
        if args.quick:
            diag.shutdown([Diagnostics.ACTION_DISPLAY_FAIL])
        else:
            diag.shutdown([Diagnostics.ACTION_DISPLAY_PASS])
        diag.disconnect()

    # diag command
    diagparser = cmdargs.add_parser(
        u'diag',
        help=u'retrieves diagnostic info; includes state changes like reboot, sleep etc'
    )
    diagcmd = diagparser.add_subparsers()

    # mobile gestalt command
    gestaltcmd = diagcmd.add_parser(
        u'gestalt',
        help=u'retrieves mobile gestalt values'
    )
    gestaltcmd.add_argument(
        u'name',
        nargs=u'?',
        help=u'the name of the gestalt value to retrieve'
    )
    gestaltcmd.set_defaults(func=cmd_gestalt)

    # all command
    allcmd = diagcmd.add_parser(
        u'all',
        help=u'retrieves all diagnostic info; gasguage, wifi and nand'
    )
    allcmd.set_defaults(func=cmd_all)

    # ioreg command
    ioregcmd = diagcmd.add_parser(
        u'ioreg',
        help=u'retrieves ioregistry information from the device; defaults to the IOService plane'
    )
    ioregcmd.add_argument(
        u'-p',
        metavar=u'plane',
        dest=u'plane',
        help=u'the ioregistry plane to retrieve; IODeviceTree, IOPower, IOService etc'
    )
    ioregcmd.add_argument(
        u'-n',
        metavar=u'name',
        dest=u'name',
        help=u'ioregistry entry name to retrieve'
    )
    ioregcmd.add_argument(
        u'-c',
        metavar=u'class',
        dest=u'cls',
        help=u'ioregistry entry class to retrieve'
    )
    ioregcmd.set_defaults(func=cmd_ioreg)

    # sleep command
    sleepcmd = diagcmd.add_parser(
        u'sleep',
        help=u'put the device to sleep'
    )
    sleepcmd.set_defaults(func=cmd_sleep)

    # restart command
    restartcmd = diagcmd.add_parser(
        u'restart',
        help=u'restarts the device'
    )
    restartcmd.add_argument(
        u'-q',
        dest=u'quick',
        action=u'store_true',
        help=u'restarts quickly - without properly shutting down'
    )
    restartcmd.set_defaults(func=cmd_restart)

    # shutdown command
    shutdowncmd = diagcmd.add_parser(
        u'shutdown',
        help=u'shutsdown the device'
    )
    shutdowncmd.add_argument(
        u'-q',
        dest=u'quick',
        action=u'store_true',
        help=u'shutsdown quickly'
    )
    shutdowncmd.set_defaults(func=cmd_shutdown)