afc.py

#!/usr/bin/env python
# -*- coding: utf8 -*-
#
# $Id$
#
# Copyright (c) 2012-2014 "dark[-at-]gotohack.org"
#
# This file is part of pymobiledevice
#
# pymobiledevice is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
#


from construct.core import Struct
from construct.lib.container import Container
from construct.macros import String, ULInt64
from lockdown import LockdownClient
import struct
from cmd import Cmd
import os
from util import hexdump, parsePlist
from pprint import pprint
import plistlib
import time 

MODEMASK =  0o0000777

AFC_OP_STATUS          = 0x00000001
AFC_OP_DATA            = 0x00000002    #Data */
AFC_OP_READ_DIR        = 0x00000003    #ReadDir */
AFC_OP_READ_FILE       = 0x00000004    #ReadFile */
AFC_OP_WRITE_FILE      = 0x00000005    #WriteFile */
AFC_OP_WRITE_PART      = 0x00000006    #WritePart */
AFC_OP_TRUNCATE        = 0x00000007    #TruncateFile */
AFC_OP_REMOVE_PATH     = 0x00000008    #RemovePath */
AFC_OP_MAKE_DIR        = 0x00000009    #MakeDir */
AFC_OP_GET_FILE_INFO   = 0x0000000a    #GetFileInfo */
AFC_OP_GET_DEVINFO     = 0x0000000b    #GetDeviceInfo */
AFC_OP_WRITE_FILE_ATOM = 0x0000000c    #WriteFileAtomic (tmp file+rename) */
AFC_OP_FILE_OPEN       = 0x0000000d    #FileRefOpen */
AFC_OP_FILE_OPEN_RES   = 0x0000000e    #FileRefOpenResult */
AFC_OP_READ            = 0x0000000f    #FileRefRead */
AFC_OP_WRITE           = 0x00000010    #FileRefWrite */
AFC_OP_FILE_SEEK       = 0x00000011    #FileRefSeek */
AFC_OP_FILE_TELL       = 0x00000012    #FileRefTell */
AFC_OP_FILE_TELL_RES   = 0x00000013    #FileRefTellResult */
AFC_OP_FILE_CLOSE      = 0x00000014    #FileRefClose */
AFC_OP_FILE_SET_SIZE   = 0x00000015    #FileRefSetFileSize (ftruncate) */
AFC_OP_GET_CON_INFO    = 0x00000016    #GetConnectionInfo */
AFC_OP_SET_CON_OPTIONS = 0x00000017    #SetConnectionOptions */
AFC_OP_RENAME_PATH     = 0x00000018    #RenamePath */
AFC_OP_SET_FS_BS       = 0x00000019    #SetFSBlockSize (0x800000) */
AFC_OP_SET_SOCKET_BS   = 0x0000001A    #SetSocketBlockSize (0x800000) */
AFC_OP_FILE_LOCK       = 0x0000001B    #FileRefLock */
AFC_OP_MAKE_LINK       = 0x0000001C    #MakeLink */
AFC_OP_SET_FILE_TIME   = 0x0000001E    #set st_mtime */

AFC_E_SUCCESS                = 0
AFC_E_UNKNOWN_ERROR          = 1
AFC_E_OP_HEADER_INVALID      = 2
AFC_E_NO_RESOURCES           = 3
AFC_E_READ_ERROR             = 4
AFC_E_WRITE_ERROR            = 5
AFC_E_UNKNOWN_PACKET_TYPE    = 6
AFC_E_INVALID_ARG            = 7
AFC_E_OBJECT_NOT_FOUND       = 8
AFC_E_OBJECT_IS_DIR          = 9
AFC_E_PERM_DENIED            =10
AFC_E_SERVICE_NOT_CONNECTED  =11
AFC_E_OP_TIMEOUT             =12
AFC_E_TOO_MUCH_DATA          =13
AFC_E_END_OF_DATA            =14
AFC_E_OP_NOT_SUPPORTED       =15
AFC_E_OBJECT_EXISTS          =16
AFC_E_OBJECT_BUSY            =17
AFC_E_NO_SPACE_LEFT          =18
AFC_E_OP_WOULD_BLOCK         =19
AFC_E_IO_ERROR               =20
AFC_E_OP_INTERRUPTED         =21
AFC_E_OP_IN_PROGRESS         =22
AFC_E_INTERNAL_ERROR         =23

AFC_E_MUX_ERROR              =30
AFC_E_NO_MEM                 =31
AFC_E_NOT_ENOUGH_DATA        =32
AFC_E_DIR_NOT_EMPTY          =33

AFC_FOPEN_RDONLY   = 0x00000001 #/**< r   O_RDONLY */
AFC_FOPEN_RW       = 0x00000002 #/**< r+  O_RDWR   | O_CREAT */
AFC_FOPEN_WRONLY   = 0x00000003 #/**< w   O_WRONLY | O_CREAT  | O_TRUNC */
AFC_FOPEN_WR       = 0x00000004 #/**< w+  O_RDWR   | O_CREAT  | O_TRUNC */
AFC_FOPEN_APPEND   = 0x00000005 #/**< a   O_WRONLY | O_APPEND | O_CREAT */
AFC_FOPEN_RDAPPEND = 0x00000006 #/**< a+  O_RDWR   | O_APPEND | O_CREAT */

AFC_HARDLINK = 1
AFC_SYMLINK = 2

AFC_LOCK_SH = 1 | 4  #/**< shared lock */
AFC_LOCK_EX = 2 | 4  #/**< exclusive lock */
AFC_LOCK_UN = 8 | 4  #/**< unlock */


AFCMAGIC = "CFA6LPAA"
AFCPacket = Struct("AFCPacket",
                   String("magic", 8,),
                   ULInt64("entire_length"),
                   ULInt64("this_length"),
                   ULInt64("packet_num"),
                   ULInt64("operation")
                   )


class AFCClient(object):
    def __init__(self, lockdown=None, serviceName="com.apple.afc", service=None):
        self.serviceName = serviceName
        if lockdown:
            self.lockdown = lockdown
        else:
            self.lockdown = LockdownClient()

        if service:
            self.service = service
        else:
            self.service = self.lockdown.startService(serviceName)
        self.packet_num = 0


    def stop_session(self):
        print "Disconecting..."
        self.service.close()
  
        
    def dispatch_packet(self, operation, data, this_length=0):
        afcpack = Container(magic=AFCMAGIC,
                   entire_length=40 + len(data),
                   this_length=40 + len(data),
                   packet_num=self.packet_num,
                   operation=operation)
        if this_length:
            afcpack.this_length = this_length
        header = AFCPacket.build(afcpack)
        self.packet_num += 1
        self.service.send(header + data)
    

    def receive_data(self):
        res = self.service.recv(40)
        status = AFC_E_SUCCESS
        data = ""
        if res:
            res = AFCPacket.parse(res)
            assert res["entire_length"] >= 40
            length = res["entire_length"] - 40
            data = self.service.recv_exact(length)
            if res.operation == AFC_OP_STATUS:
                if length != 8:
                    print "Status length != 8"
                status = struct.unpack("<Q", data[:8])[0]
            elif res.operation != AFC_OP_DATA:
                pass#print "error ?", res
        return status, data


    def do_operation(self, opcode, data=""):
        try:
            self.dispatch_packet(opcode, data)
            return self.receive_data()
        except:
            self.lockdown = LockdownClient()
            self.service = self.lockdown.startService(self.serviceName)
            return  self.do_operation(opcode, data)
    

    def list_to_dict(self, d):
        t = d.split("\x00")
        t = t[:-1]
        
        assert len(t) % 2 == 0
        res = {}
        for i in xrange(len(t)/2):
            res[t[i*2]] = t[i*2 + 1]
        return res
        

    def get_device_infos(self):
        status, infos = self.do_operation(AFC_OP_GET_DEVINFO)
        if status == AFC_E_SUCCESS:
            return self.list_to_dict(infos)


    def read_directory(self, dirname):
        status, data = self.do_operation(AFC_OP_READ_DIR, dirname)
        if status == AFC_E_SUCCESS:
            return filter(lambda x:x!="", data.split("\x00"))
        return []
    

    def make_directory(self, dirname):
        status, data = self.do_operation(AFC_OP_MAKE_DIR, dirname)
        return status
    

    def remove_directory(self, dirname):
        info = self.get_file_info(dirname)
        if not info or info.get("st_ifmt") != "S_IFDIR":
            print "remove_directory: %s not S_IFDIR" % dirname
            return
        
        for d in self.read_directory(dirname):
            if d == "." or d == ".." or d == "":
                continue
            
            info = self.get_file_info(dirname + "/" + d)
            if info.get("st_ifmt") == "S_IFDIR":
                self.remove_directory(dirname + "/" + d)
            else:
                print dirname + "/" + d
                self.file_remove(dirname + "/" + d)
        assert len(self.read_directory(dirname)) == 2 #.. et .
        return self.file_remove(dirname)
        

    def get_file_info(self, filename):
        status, data = self.do_operation(AFC_OP_GET_FILE_INFO, filename)
        if status == AFC_E_SUCCESS:
            return self.list_to_dict(data)


    def make_link(self, target, linkname, type=AFC_SYMLINK):
        status, data = self.do_operation(AFC_OP_MAKE_LINK, struct.pack("<Q", type) + target + "\x00" + linkname + "\x00")
        print "make_link", status
        return status
        

    def file_open(self, filename, mode=AFC_FOPEN_RDONLY):
        status, data = self.do_operation(AFC_OP_FILE_OPEN, struct.pack("<Q", mode) + filename + "\x00")
        if data:
            handle = struct.unpack("<Q", data)[0]
            return handle
    

    def file_close(self, handle):
        status, data = self.do_operation(AFC_OP_FILE_CLOSE, struct.pack("<Q", handle))
        return status
    

    def file_remove(self, filename):
        status, data = self.do_operation(AFC_OP_REMOVE_PATH, filename + "\x00")
        return status 
    

    def file_rename(self, old, new):
        status, data = self.do_operation(AFC_OP_RENAME_PATH, old + "\x00" + new + "\x00")
        return status
    

    def file_read(self, handle, sz):
        MAXIMUM_READ_SIZE = 1 << 16
        data = ""
        while sz > 0:
            if sz > MAXIMUM_READ_SIZE:
                toRead = MAXIMUM_READ_SIZE
            else:
                toRead = sz
	    try:
		self.dispatch_packet(AFC_OP_READ, struct.pack("<QQ", handle, toRead))
		s, d = self.receive_data()
	    except:
		self.lockdown = LockdownClient()
		self.service = self.lockdown.startService("com.apple.afc")
		return  self.file_read(handle, sz)

            if s != AFC_E_SUCCESS:
                break
            sz -= toRead
            data += d
        return data


    def file_write(self, handle, data):
        MAXIMUM_WRITE_SIZE = 1 << 15
        hh = struct.pack("<Q", handle)
        segments = len(data) / MAXIMUM_WRITE_SIZE
	try:
	    for i in xrange(segments):
		self.dispatch_packet(AFC_OP_WRITE,
		                 hh + data[i*MAXIMUM_WRITE_SIZE:(i+1)*MAXIMUM_WRITE_SIZE],
			             this_length=48)
		s, d = self.receive_data()
		if s != AFC_E_SUCCESS:
		    print "file_write error %d" % s
		    break
	    if len(data) % MAXIMUM_WRITE_SIZE:
		self.dispatch_packet(AFC_OP_WRITE,
		                     hh + data[segments*MAXIMUM_WRITE_SIZE:],
			             this_length=48)
		s, d = self.receive_data()
	except:
	    self.lockdown = LockdownClient()
	    self.service = lockdown.startService(serviceName)
	    self.file_write(handle,data)
        return s
    

    def get_file_contents(self, filename):
        info = self.get_file_info(filename)
        if info:
            if info['st_ifmt'] == 'S_IFLNK':
                filename =  info['LinkTarget']

            if info['st_ifmt'] == 'S_IFDIR':
                print "%s is directory..." % filename
                return
	    
            print "Reading %s" % filename
            h = self.file_open(filename)
            if not h:
                return
            d = self.file_read(h, int(info["st_size"]))
            self.file_close(h)
            return d
	    return


    def set_file_contents(self, filename, data):
        h = self.file_open(filename, AFC_FOPEN_WR)
        if not h:
            return
        d = self.file_write(h, data)
        self.file_close(h)

    def dir_walk(self,dir,file_list=[]):
        d = os.path.abspath(dir)
        file_list = []
        for file in [file for file in self.read_directory(d) if not file in [".",".."]]:
            path = os.path.join(d,file)
            info =  self.get_file_info(path)
            if info:
                if info['st_ifmt'] == 'S_IFDIR':
                    file_list += self.dir_walk(path,file_list)
                    info['path'] = path
                    file_list.append(info)
	    return file_list

class AFCShell(Cmd):

    def __init__(self, completekey='tab', stdin=None, stdout=None, client=None):

        Cmd.__init__(self, completekey=completekey, stdin=stdin, stdout=stdout)
        self.lockdown = LockdownClient()
        if client:
            self.afc = client
        else:
            self.afc = AFCClient(self.lockdown, "com.apple.afc")
        self.curdir = "/"
        self.prompt = "AFC$ " + self.curdir + " "
        self.complete_cat = self._complete
        self.complete_ls = self._complete
    

    def do_exit(self, p):
        return True
    

    def do_quit(self, p):
        return True
    

    def do_pwd(self, p):
        print self.curdir
    
            
    def do_link(self, p):

        z = p.split()
        self.afc.make_link(AFC_SYMLINK, z[0], z[1])
    
        
    def do_cd(self, p):

        if not p.startswith("/"):
            new = self.curdir + "/" + p
        else:
            new = p
        
        new = os.path.normpath(new).replace("\\","/").replace("//","/")    
        d = self.afc.read_directory(new)
        if d:
            self.curdir = new
            self.prompt = "AFC$ %s " % new
        else:
            print "%s does not exists" % new
    
   
    def _complete(self, text, line, begidx, endidx):
        filename = text.split("/")[-1]
        dirname = "/".join(text.split("/")[:-1])
        return [dirname + "/" + x for x in self.afc.read_directory(self.curdir + "/" + dirname) if x.startswith(filename)]
    

    def do_ls(self, p):

        d = self.afc.read_directory(self.curdir + "/" + p)
        if d:
            for dd in d:
                print dd
    

    def do_cat(self, p):

        data = self.afc.get_file_contents(self.curdir + "/" + p)
        if data and p.endswith(".plist"):
            pprint(parsePlist(data))
        else:
            print data


    def do_rm(self, p):
        f =  self.afc.get_file_info(self.curdir + "/" + p)
        if f['st_ifmt'] == 'S_IFDIR':
            d = self.afc.remove_directory(self.curdir + "/" + p)
        else:
            d = self.afc.file_remove(self.curdir + "/" + p)
    
        
    def do_pull(self, user_args):
        
        args = user_args.split()
        if len(args) != 2:
            out = "."
            path = user_args
        else: 
            out = args[1]
            path = args[0]
            
        
        f =  self.afc.get_file_info(path)
        if not f:
            print "Source file does not exist.."
            return

        out_path = out + "/" + path
        if f['st_ifmt'] == 'S_IFDIR':

            if not os.path.isdir(out_path):
                os.makedirs(out_path, MODEMASK)
            
            for d in self.afc.read_directory(path):
                if d == "." or d == ".." or d == "":
                    continue
                self.do_pull(path + "/" + d + " " + out)
        else:
            data = self.afc.get_file_contents(path)
            if data:
                if data and path.endswith(".plist"):
                    z = parsePlist(data)
                    plistlib.writePlist(z, out_path)
                else:
                    with open(out_path, 'wb+') as f:
                        f.write(data)                

    def do_push(self, p):

        fromTo = p.split()
        if len(fromTo) != 2:
            return
        print "from %s to %s" % (fromTo[0], fromTo[1])
        if os.path.isdir(fromTo[0]):
            self.afc.make_directory(os.path.join(fromTo[1]))
            for x in os.listdir(fromTo[0]):
                if x.startswith("."):
                    continue
                path = os.path.join(fromTo[0],x)
                self.do_push(path + " " + fromTo[1]+ "/" + path)
        else:
            if not fromTo[0].startswith("."):
                data = open(fromTo[0], "rb").read()
                self.afc.set_file_contents(self.curdir + "/" + fromTo[1], data)

            
    def do_head(self, p):
        print self.afc.get_file_contents(self.curdir + "/" + p)[:32]


    def do_hexdump(self, p):

        t = p.split(" ")
        l = 0
        if len(t) < 1:
            return
        if len(t) == 2:
            l = int(t[1])
        z = self.afc.get_file_contents(self.curdir + "/" + t[0])
        if not z:
            return
        if l:
            z = z[:l]
        hexdump(z)

    
    def do_mkdir(self, p):
        print self.afc.make_directory(p)


    def do_infos(self, p):
        print self.afc.get_device_infos()
        

if __name__ == "__main__":
    AFCShell().cmdloop("Hello iPhone!")