# Very alpha still, but getting there...

# Yeah, I like it this way
require 'pp'
require 'set'
require 'zlib'
require 'base64'
require 'socket'
require 'openssl'
require 'stringio'
require 'net/http'
require 'net/https'
require 'osx/cocoa'
require 'rexml/document'

begin
  require 'usb' 
  $usb_available = true
rescue 
  $usb_available = false
end

Thread.abort_on_exception = true

class String
  def hexdump
    buffer = ""
    (0..length).step(16) do |i|
      first_half = self[i, 8].unpack("C*").map{|x| "%02x" % x}.join(" ")
      second_half = (self[i + 8, 8] || "").unpack("C*").map{|x| "%02x" % x}.join(" ")
      first_half += " " * (8 * 2 + (8 - 1) - first_half.length)
      second_half += " " * (8 * 2 + (8 - 1) - second_half.length)
      buffer += first_half + "  " + second_half + "\t" + self[i, 16].unpack("C*").map{|x| ((32..128).include?(x) ? x : ?.).chr}.join + "\n"
    end
    puts buffer
  end
  
  def printable?
    self.unpack("C*").all?{|x| ((32..128).entries + [9, 10]).include?(x)}
  end
  
  def to_hex
    self.unpack("C*").map{|x| "%02x" % x}.join
  end
end

class Integer
  @@host_is_little = "\x01\x02\x03\x04".unpack("N4")[0] == 0x01020304
  
  def byteswap
    value = 0
    self.size.times do |i|
      value |= (self & (0xFF << (8 * i))) << (8 * ((self.size - 1) - (2 * i)))
    end
    value
  end
  
  def big_to_host
    @@host_is_little ? self.byteswap : self
  end
  
  def little_to_host
    @@host_is_little ? self : self.byteswap
  end
end

class PList
  attr_reader :object
  
  def initialize(data)
    @object = data.printable? ? decode_xml_plist(data) : decode_binary_plist(data) 
  end
  
  def get_int(size)
    value = 0
    size.times{value = (value << 8) | @io.read(1).unpack("C")[0]}
    value
  end
  
  def check_int
    marker = @io.read(1).unpack("C")[0]
    raise "Int expected" unless marker & 0xf0 == 0x10
    get_int(1 << (marker & 0xf))
  end
  
  def read_object
    marker = @io.read(1).unpack("C")[0]
    
    case marker & 0xf0
    when 0x00 # Null
      case marker & 0x0f
      when 0: nil
      when 8: false
      when 9: true
      when 15: "<fill>"
      end
    when 0x10 # Int
      get_int(1 << marker & 0x0f)
    when 0x20: "**** Real ****"
    when 0x30: "**** Date ****"
    when 0x40, 0x50 # Data / ASCII
      size = marker & 0xf
      size = check_int if size == 15
      
      @io.read(size)
    when 0x60 # Unicode16
    when 0x80 # UID
    when 0xa0 # Array
      size = marker & 0x0f
      size = check_int if size == 15
      
      value_indices = []
      size.times {value_indices << get_int(@offset_int_size)}
      
      value_indices.map {|i| @io.seek(@offset_table[i]); read_object}
    when 0xd0 # dict
      size = marker & 0xf
      size = check_int if size == 15
      
      key_indices = []
      value_indices = []
      
      size.times {key_indices << get_int(@offset_int_size)}
      size.times {value_indices << get_int(@offset_int_size)}
      
      keys = key_indices.map{|i| @io.seek(@offset_table[i]); read_object}
      values = value_indices.map{|i| @io.seek(@offset_table[i]); read_object}
            
      keys.zip(values).inject({}){|accum, pair| accum[pair[0]] = pair[1]; accum}
    end    
  end
  
  def decode_binary_plist(data)
    @io = StringIO.new(data)
    raise "Invalid binary plist" unless @io.read(6) == "bplist" && (0..2).member?(@io.read(2).to_i )
        
    @io.seek(-32, IO::SEEK_END)
    unused, @offset_int_size, @object_ref_size, @num_objects, @top_object, @offset_table_offset = @io.read(32).unpack("a6CCQQQ")
    
    @num_objects = @num_objects.big_to_host
    @top_object = @top_object.big_to_host
    @offset_table_offset = @offset_table_offset.big_to_host
    
    @io.seek(@offset_table_offset)
    
    @offset_table = []
    @num_objects.times{@offset_table << get_int(@offset_int_size)}
    
    @io.seek(@offset_table[@top_object])
    
    read_object
  end

  def decode_xml_element(element)
    return decode_xml_element(element.next_sibling_node) if element.class == REXML::Text && element.value.strip.empty?
    
    case element.name
    when "plist": decode_xml_element(element.children[0])
    when "dict"
      hash = {}
      element.each_child do |child|
        if child.kind_of?(REXML::Element) && child.name == "key"
          hash[child.children[0].value] = decode_xml_element(child.next_sibling_node)
        end
      end
      hash
    when "array"
      array = []
      element.each_child{|child| array << decode_xml_element(child) if child.kind_of?(REXML::Element)}
      array
    when "integer": element.children[0].value.to_i
    when "real": element.children[0].value.to_f
    when "string": element.children[0].value
    when "data": Base64.decode64(element.children[0].value)
    when "true": true
    when "false": false
    end
  end

  def decode_xml_plist(data)
    decode_xml_element(REXML::Document.new(data).root)
  end
end

class CPIO
  def initialize(input)
    @io = input.kind_of?(String) ? StringIO.new(input) : input
        
    while true do
      magic, dev, ino, mode, uid, gid, nlink, rdev, mtime, namesize, filesize = @io.read(76).unpack("a6a6a6a6a6a6a6a6a11a6a11")
    
      # Weird format uses octal ascii? 
      file_name = @io.read(namesize.to_i(8)).chop
      file_data = @io.read(filesize.to_i(8))
            
      break if file_name == "TRAILER!!!"
      
      puts file_name
    end
  end
end

class SocketRelay
  def self.[](a, b)
    self.new(a, b)
  end
  
  def initialize(a, b)
    @a, @b = a, b
    
    # TODO: This should be done with a select
    
    @a_to_b = Thread.new do
      while true
        value = @a.readpartial(4096) rescue break
        p value
        @b.write(value)
      end
    end
    
    @b_to_a = Thread.new do
      while true
        value = @b.readpartial(4096) rescue break
        p value
        @a.write(value)
      end
    end
  end
end

class USBTCPSocket
  attr_reader :device_id, :product_id, :serial_no
  
  def send_packet(socket, packet_type, data)
    packet = [data.length + 16, 0, packet_type, @session_number].pack("V4") + data
    socket.write(packet)
  end

  def recv_packet(socket)
    header = socket.read(16)
    packet_length, unk, packet_type, tag = header.unpack("V4")
    data = socket.read(packet_length - 16)
    [packet_type, tag, data]
  end
  
  def initialize(port, override_device_id = nil)
    @session_number = 1
    
    if override_device_id
      @device_id = override_device_id
    else
      # This just grabs the device id and serial number. Once you've run the hello protocol it won't accept anything else
      socket = UNIXSocket.new("/var/run/usbmuxd")

      puts "Sending hello packet"
      send_packet(socket, 3, "")
      p recv_packet(socket)
      @device_id, @product_id, @serial_no = recv_packet(socket)[2].unpack("Vva40")    
    end
    
    puts "Device ID: 0x#{@device_id.to_s(16)}"
    
    @use_ssl = false
    
    done = false
    
    until done do
      @socket = UNIXSocket.new("/var/run/usbmuxd")
      puts "Retrying connection to port #{port}..."
      send_packet(@socket, 2, [device_id, port, 0].pack("Vnn"))
      response = recv_packet(@socket)[2]
      p response
      done = response.unpack("N")[0] == 0
      @socket.close unless done
      @session_number += 1
    end
    
    puts "Connected to port #{port}"
  end
  
  def use_ssl=(value)
    if value && !@use_ssl
      @use_ssl = true
      @plain_socket = @socket
      @socket = OpenSSL::SSL::SSLSocket.new(@plain_socket, OpenSSL::SSL::SSLContext.new(:TLSv1))
      @socket.connect
    elsif !value && @use_ssl
      @use_ssl = false
      @socket.close
      @socket = @plain_socket
    end
  end
  
  def write(data)
    @socket.write(data)
  end
  
  def read(size)
    @socket.read(size)
  end
  
  def gets(separator = $/)
    @socket.gets(separator)
  end
  
  def readpartial(size)
    @socket.readpartial(size)
  end
  
  def close
    @socket.close
  end
end

class Service
  def self.handle(type)
    (@@handlers ||= Hash.new{|hash, key| hash[key] = USBTCPSocket})[type] = self
  end
  
  def self.[](type)
    @@handlers[type]
  end
  
  def self.all
    @@handlers.keys
  end
end

class PropertyListService < Service
  def initialize(port, device_id = nil)
    @socket = USBTCPSocket.new(port, device_id)
  end
  
  def request_plist(data, format = OSX::NSPropertyListXMLFormat_v1_0)
    write_plist(data, format)
    read_plist
  end
  
  def write_plist(data, format = OSX::NSPropertyListXMLFormat_v1_0)
    if data.kind_of?(Hash) || data.kind_of?(Array)
      data = data.to_plist
    end
    
    @socket.write([data.length].pack("N") + data)
  end
  
  def read_plist
    size = @socket.read(4).unpack("N")[0]
    buffer = @socket.read(size)
    OSX.load_plist(buffer)
  end
end

class PortForwarderService < Service
  def initialize(port)
    super(port)
    
    # TODO: fill me in
  end
end

class LockdownService < PropertyListService
  handle "com.apple.mobile.lockdown"
  
  def initialize(port, label = "iTunes", host_id = nil, pair_record_path = nil)
    super(port)
    
    @label = label
    
    pair_record_path ||= "#{File.expand_path("~")}/Library/Lockdown/#{@socket.serial_no}.plist"
    
    raise "QueryType request failed" unless query_type["Result"] == "Success"
    raise "Not talking to lockdownd! #{query_type["Type"]} instead!" unless query_type["Type"] == "com.apple.mobile.lockdown"
    
    pair_record = Hash[*OSX.load_plist(File.read(pair_record_path)).select{|key, value| %w(DeviceCertificate HostCertificate HostID RootCertificate).include?(key)}.map{|key, value| [key, value.respond_to?(:rubyString) ? value.rubyString : value]}.flatten]
  
    raise "Unable to validate pairing records" unless validate_pair(pair_record)["Result"] == "Success"
    
    raise "Unable to start session" unless start_session(pair_record["HostID"])["Result"] == "Success"
  end
  
  def request(type, hash = {})
    request_plist({"Label" => @label, "Request" => type}.merge(hash))
  end
  
  # The get value calls seem to return meaningful values for the following domains:
  #   com.apple.mobile.lockdown
  #   com.apple.mobile.iTunes
  #   com.apple.mobile.battery
  #   com.apple.springboard.curvedBatteryCapacity
  #   com.apple.mobile.internal
  #   com.apple.mobile.debug
  #   com.apple.mobile.restriction
  #   com.apple.mobile.sync_data_class
  #   com.apple.mobile.data_sync
  #   com.apple.mobile.nikita
  #   com.apple.fairplay
  #   com.apple.international
  #   com.apple.disk_usage
  #   and after xcode has done stuff: com.apple.xcode.developerdomain
  
  def get_value(key, domain = "")
    request("GetValue", {"Key" => key}.merge(domain == "" ? {} : {"Domain" => domain}))["Value"]
  end
  
  def pair(root_certificate, device_certificate, host_certificate, host_id)
    request("Pair", "PairRecord" => {
      "DeviceCertificate" => device_certificate,
      "HostCertificate" => host_certificate,
      "HostID" => host_id,
      "RootCertificate" => root_certificate
    })
  end
  
  def unpair
    request("Unpair")
  end
  
  def activate(activation_record)
    request("Activate", "ActivationRecord" => activation_record)
  end
  
  def deactivate
    request("Deactivate")
  end
  
  def get_values(domain = "")
    # TODO: check for success
    request("GetValue", domain == "" ? {} : {"Domain" => domain})["Value"]
  end
  
  def query_type
    # Might as well cache this one
    @query_type ||= request("QueryType")
  end
  
  def validate_pair(pair_record)
    request("ValidatePair", "PairRecord" => pair_record)
  end
  
  def enter_recovery
    request("EnterRecovery")
  end
  
  def start_session(host_id)
    response = request("StartSession", "HostID" => host_id)
    
    @socket.use_ssl = response["EnableSessionSSL"]
    
    response
  end
  
  def start_service(service_name, handler = nil)
    response = request("StartService", "Service" => service_name)
    
    raise "Unable to start service #{service_name}" unless response["Result"] == "Success"
    
    (handler || Service[service_name]).new(response["Port"])
  end
end

class DeviceLinkService < PropertyListService
  def initialize(port)
    super(port)

    version_exchange = read_plist
    raise "Version exchange message expected" unless version_exchange[0] == "DLMessageVersionExchange"
    
    response = request_plist(["DLMessageVersionExchange", "DLVersionsOk"], OSX::NSPropertyListBinaryFormat_v1_0)
        
    raise "Device not ready" unless response[0] == "DLMessageDeviceReady"
  end
  
  def close
    write_plist(["DLMessageDisconnect", "All good things must come to an end."])
  end
end

class ScreenShotService < DeviceLinkService
  handle "com.apple.mobile.screenshotr"

  def take_screenshot
    response = request_plist(["DLMessageProcessMessage", {"MessageType" => "ScreenShotRequest"}])
    
    raise "Invalid reply type" unless response[0] == "DLMessageProcessMessage"
    raise "Invalid reply message type: #{response[1]["MessageType"]}" unless response[1]["MessageType"] == "ScreenShotReply"
    
    response[1]["ScreenShotData"].rubyString
  end
end

class CrashReportService < DeviceLinkService
  handle "com.apple.crashreportcopy"
  
  def system_info
    request_plist(["DLMessageProcessMessage", {"VersionNumber" => 101, "CrashReportRequest" => "GetSysInfo"}])
  end
end

class ImageMounterService < PropertyListService
  handle "com.apple.mobile.mobile_image_mounter"
  
  def lookup_image(image_type)
    request("LookupImage", "ImageType" => "Developer")
  end
  
  def mount_image(path, signature, type)
    response = request("MountImage", "ImagePath" => path, "ImageSignature" => signature, "ImageType" => type)
  end
  
  def close
    request("Hangup")
  end
  
  def request(command, hash = {})
    request_plist({"Command" => command}.merge(hash))
  end
end

class NotificationProxyService < PropertyListService
  handle "com.apple.mobile.notification_proxy"
  
  attr_reader :listener
  
  def initialize(port)
    super(port)
    
    @observers = Hash.new{|hash, key| hash[key] = Set[]}
    
    @listener = Thread.start do
      while true
        plist = read_plist
        
        if plist["Command"] == "RelayNotification"
          @observers[plist["Name"]].each do |handler|
            handler.call
          end
        end
      end
    end
  end
  
  # Valid notifications include:
  #   com.apple.language.changed
  #   com.apple.AddressBook.PreferenceChanged
  #   com.apple.mobile.data_sync.domain_changed
  #   com.apple.mobile.lockdown.device_name_changed
  #   com.apple.mobile.developer_image_mounted
  #   com.apple.mobile.lockdown.trusted_host_attached
  #   com.apple.mobile.lockdown.host_detached
  #   com.apple.mobile.lockdown.host_attached
  #   com.apple.mobile.lockdown.phone_number_changed
  #   com.apple.mobile.lockdown.registration_failed
  #   com.apple.mobile.lockdown.activation_state
  #   com.apple.mobile.lockdown.brick_state
  #   com.apple.itunes-client.syncCancelRequest
  #   com.apple.itunes-client.syncSuspendRequest
  #   com.apple.itunes-client.syncResumeRequest
  #   com.apple.springboard.attemptactivation
  #   com.apple.mobile.application_installed
  #   com.apple.mobile.application_uninstalled
  
  def observe(name, &block)
    write_plist("Command" => "ObserveNotification", "Name" => name)
    
    @observers[name] << block
  end
  
  def post(name)
    write_plist("Command" => "PostNotification", "Name" => name)
  end
end

class InstallationProxyService < PropertyListService
  handle "com.apple.mobile.installation_proxy"
  
  def browse(options = {})
    result = request("Browse", options)
    
    done = false
    
    until done do
      plist = read_plist
      done = plist["Status"] == "Complete"
    end
    
    # TODO: make it work better
    result
  end
  
  def install(package_path, type = "User", &progress_callback)
    request("Install", "PackagePath" => package_path, "ClientOptions" => {"PackageType" => type})
    
    done = false
    
    until done do
      plist = read_plist
      done = plist["Status"] == "Complete"
    end
  end
  
  def request(command, hash = {})
    request_plist({"Command" => command}.merge(hash))
  end
end

class SyslogRelayService < Service
  handle "com.apple.syslog_relay"
  
  attr_reader :listener
  
  def initialize(port)
    @socket = USBTCPSocket.new(port)
    
    @syslog = ""
    @observers = Set[]
    
    @listener = Thread.new do
      while true
        # syslog_relay will send at most 16383 bytes at a time
        buffer = @socket.readpartial(16383)
        
        @observers.each{|observer| observer.call(buffer)}
        
        @syslog << buffer
      end
    end
  end
  
  def observe(&block)
    @observers << block
  end
end

class MISService < PropertyListService
  handle "com.apple.misagent"
  
  def initialize(port)
    super(port)
  end
  
  # Only provisioning profiles seem to be supported in misagent for now (as of 5a345)
  def install_profile(profile, profile_type = "Provisioning")
    request("Install", "Profile" => profile, "ProfileType" => profile_type)
  end
  
  # Fixme: What does this do?
  def copy_profile(profile_type = "Provisioning")
    request("Copy", "ProfileType" => profile_type)
  end
  
  def remove_profile(profile_id, profile_type = "Provisioning")
    request("Remove", "ProfileID" => profile_id, "ProfileType" => profile_type)
  end
  
  def request(message_type, hash = {})
    request_plist({"MessageType" => message_type}.merge(hash))
  end
end

class ApplistService < Service
  handle "com.apple.debugserver.applist"
  
  def initialize(port)
    @socket = USBTCPSocket.new(port)
  end
  
  def applist
    # Fixme: Hardcoded
    OSX.load_plist(@socket.readpartial(128 * 1024))
  end
end

class FileRelayService < PropertyListService
  handle "com.apple.mobile.file_relay"
  
  # The following are valid sources
  #   AppleSupport
  #   Network
  #   VPN
  #   WiFi
  #   UserDatabases
  #   CrashReporter
  #   tmp
  #   SystemConfiguration
  
  def request_source(*sources)
    result = request_plist("Sources" => sources.flatten)
    
    raise "Unable to retrieve sources" unless result["Status"] == "Acknowledged"
    
    gzip_reader = Zlib::GzipReader.new(@socket)
    CPIO.new(gzip_reader)
  end
end

class BackupService < DeviceLinkService
  handle "com.apple.mobilebackup"
  
  def initialize(port)
    super(port)
    
    # TODO: finish me
    ["DLMessageProcessMessage", {"BackupMessageTypeKey" => "BackupMessageBackupRequest", 
                                 "BackupComputerBasePathKey"=>"/",
                                 "BackupManifestKey" =>
                                  {"Data" => "moo",
                                   "AuthVersion" => "1.0",
                                   "AuthData" => "Forty Two",
                                   "Auth Signature" => "SHA1"}}]
  end
end

class AFCService < Service
  handle "com.apple.afc"
  handle "com.apple.afc2"
  
  class AFCIO
    def write(data)
      
    end
    
    def read(size)
      
    end
  end
  
  def initialize(port)
    @socket = USBTCPSocket.new(port)
    
    @sequence_number = 0
  end
  
  def send_frame(type, data = "", header_size = data.length + 40)
    frame = "CFA6LPAA" + [data.length + 40, 0, header_size, 0, @sequence_number += 1, 0, type, 0].pack("V*") + data
    frame.hexdump
    @socket.write(frame)
  end
  
  def recv_frame()
    magic, size = @socket.read(16).unpack("a8V*")
    
    raise "Invalid frame" unless magic == "CFA6LPAA"
    
    header_size = @socket.read(8).unpack("V")[0] # Ignoring other 4 bytes
    sequence_number = @socket.read(8).unpack("V")[0] # Ignoring other 4 bytes

    header_rest = @socket.read(header_size - 32)
    
    @socket.read(size - header_size)
  end
  
  def sysinfo()
    send_frame(11)
    Hash[*recv_frame().split("\x00").map{|x| x.to_i.to_s == x ? x.to_i : x}]
  end
  
  def stat(path)
    send_frame(10, path + "\x00")
    Hash[*recv_frame().split("\x00").map{|x| x.to_i.to_s == x ? x.to_i : x}]
  end
  
  def mkdir(path)
    send_frame(9, path + "\x00")
    recv_frame()
  end
  
  def ls(path)
    send_frame(3, path + "\x00")  
    recv_frame.split("\x00")
  end
  
  # 2 appears to be readable, 3 seems to create the file
  def open(path, mode = 2)
    send_frame(13, [mode, 0].pack("V*") + path + "\x00")
    recv_frame
  end
  
  def read(file_size)
    send_frame(15, [1, 0, file_size, 0].pack("V*"))    
    recv_frame
  end
  
  def write(data)
    send_frame(16, [1, 0].pack("V*") + data, 48)
  end
end

class RestoreService < PropertyListService
  handle "com.apple.mobile.restored"
  
  def reboot
    request("Reboot")
  end
  
  def start_restore(progress_callback = nil, &data_request_handler)
    write_plist("Request" => "StartRestore", "RestoreProtocolVersion" => 11)
    p "wrote plist"
    while plist = read_plist do
      p "got plist"
      p plist
      
      if plist["MsgType"] =="DataRequestMsg"
        response = data_request_handler.call(plist["DataType"])
        write_plist(response) if response
      elsif progress_callback && plist["MsgType"] == "ProgressMsg" 
        progress_callback.call(plist["Operation"], plist["Progress"])
      elsif plist["MsgType"] == "StatusMsg"
        puts "Got status message: #{plist.inspect}"
        break if plist["Status"] == 0
      end
    end
  end
  
  def goodbye
    request("Goodbye")
  end
  
  # Valid keys (key cannot be empty):
  #   SerialNumber
  #   IMEI
  #   HardwareModel
  def query_value(key)
    request("QueryValue", "QueryKey" => key)
  end
  
  def query_type
    request("QueryType")
  end
  
  def request(command, hash = {})
    request_plist({"Request" => command}.merge(hash))
  end
end

class ASRService < Service
  def initialize(port, input)
    @socket = USBTCPSocket.new(port)
    
    if input.kind_of?(File)
      @io = input
      @size = input.stat.size
    elsif input.kind_of?(String)
      @io = StringIO.new(input)
      @size = input.size
    end
    
    raise "Unexpected command" unless read_plist["Command"] == "Initiate"
  end
  
  def start
    write_plist({
      "FEC Slice Stride" => 40,
      "Packet Payload Size" => 1450,
      "Packets Per FEC" => 25,
      "Payload" => {
        "Port" => 1,
        "Size" => @size
      },
      "Stream ID" => 1,
      "Version" => 1
    })
    
    while plist = read_plist do
      if plist["Command"] == "OOBData"
        size = plist["OOB Length"]
        offset = plist["OOB Offset"]
        
        puts "Sending #{size} OOB bytes from offset #{offset}"
        
        @io.seek(offset)
        @socket.write(@io.read(size))
      elsif plist["Command"] == "Payload"
        puts "Sending payload"
        @io.seek(0)
        
        index = 0
        
        while buffer = @io.read(0x10000) do
          @socket.write(buffer)
          index += 1
          
          if index % 16 == 0
            puts "#{index.to_f / (@size / 0x10000) * 100}% done"
          end
        end
        break
      else
        puts "Unknown ASR command #{plist.inspect}"
      end
    end
  end
  
  def read_plist
    buffer = ""
    
    while read_buffer = @socket.gets do
      puts "Read: #{read_buffer.inspect}"
      buffer << read_buffer
      break if read_buffer =~ /<\/plist>/
    end
    
    OSX.load_plist(buffer)
  end
  
  def write_plist(hash)
    @socket.write(hash.to_plist)
  end
end

class AppleDevice
  @@matches = {}
  
  def self.match(value)
    [value[:vendor_id]].flatten.each do |vendor_id|
      (@@matches ||= {})[vendor_id] ||= {}

      [value[:product_id]].flatten.each do |product_id|
        @@matches[vendor_id][product_id] = self
      end
    end
  end
  
  def self.[](vendor_id, product_id)
    @@matches[vendor_id][product_id]
  end
  
  def self.available_devices
    available_ids = {}
    @@matches.each_pair do |vendor_id, value|
      value.each_pair do |product_id, klass|
        available_ids[[vendor_id, product_id]] = klass
      end
    end
        
    USB.devices.map do |device|
      available_ids[[device.idVendor, device.idProduct]].new(device) if available_ids.has_key?([device.idVendor, device.idProduct])
    end.compact
  end
  
  def initialize(device)
    @device = device
  end
end

class NormalMode < AppleDevice
  match :vendor_id => 0x5ac, :product_id => [0x1290, 0x1291, 0x1292] # iPhone, iPod Touch, iPhone3G
  
  attr_reader :service
  
  def initialize(device)
    super(device)
  end
  
  def open
    service = PropertyListService.new(62078)

    @service = Service[service.request_plist("Request" => "QueryType")["Type"]].new(62078)
  end
end

class RecoveryV1Mode < AppleDevice
  match :vendor_id => 0x5ac, :product_id => 0x1280
end

class RecoveryV2Mode < AppleDevice
  match :vendor_id => 0x5ac, :product_id => 0x1281
  
  def initialize(device)
    super(device)
  end
  
  def open
    @handle = @device.open
    
    @handle.set_configuration(@device.configurations[0])
    @handle.claim_interface(1)
    @handle.set_altinterface(1)
    
    @receive_buffer = "\x00" * 0x1000
    
    self
  end
  
  def recv_buffer   
    size = @handle.usb_interrupt_read(0x81, @receive_buffer, 0)
    
    @receive_buffer[0, size]
  end
  
  def send_command(command)
    @handle.usb_interrupt_write(0x02, command + "\n", 0)
  end
  
  def send_file(file)
    @handle.usb_control_msg(0x21, 6, 1, 0, "", 0)
    
    index = 0
    size = 0
    while buffer = file.read(0x1000) do
      size += buffer.size
      @handle.usb_control_msg(0x21, 1, index, 0, buffer, 0)
      index += 1
    end
    
    @handle.usb_control_msg(0x21, 1, 1, 0, "", 0)
    
    buffer = "\x00" * 6
    
    status = @handle.usb_control_msg(0xa1, 3, 0, 0, buffer, 0)
    buffer.hexdump
    raise "Bad status" unless buffer[4] == 0x6
    
    status = @handle.usb_control_msg(0xa1, 3, 0, 0, buffer, 0)
    buffer.hexdump
    raise "Bad status" unless buffer[4] == 0x7
    
    status = @handle.usb_control_msg(0xa1, 3, 0, 0, buffer, 0)
    buffer.hexdump
    raise "Bad status" unless buffer[4] == 0x8
  end
  
  def close
    @handle.release_interface(1)
    @handle.usb_close
  end
end

class DFUV2Mode < AppleDevice
  match :vendor_id => 0x05ac, :product_id => [0x1222, 0x1227]  
end

#pp Service.all

x = nil

p "finding devices"

def pair
  devs = AppleDevice.available_devices
  
  if devs[0].kind_of?(NormalMode)
    x = devs[0].open

    root_ca_cert = OpenSSL::X509::Certificate.new
    root_private_key = OpenSSL::PKey::RSA.new(1024)
    root_ca_cert.public_key = root_private_key.public_key

    device_cert = OpenSSL::X509::Certificate.new
    device_cert.public_key = OpenSSL::PKey::RSA.new(x.get_value("DevicePublicKey").rubyString)  
    device_cert.sign(root_private_key, OpenSSL::Digest::Digest.new("SHA1"))
  
    host_private_key = OpenSSL::PKey::RSA.new(1024)
    host_cert = OpenSSL::X509::Certificate.new
    host_cert.public_key = host_private_key.public_key
    host_cert.sign(root_private_key, OpenSSL::Digest::Digest.new("SHA1"))

    root_pem = root_ca_cert.to_pem
    device_pem = device_cert.to_pem
    host_pem = host_cert.to_pem
  
    root_data = OSX::NSData.alloc.initWithBytes_length(root_pem, root_pem.size)
    host_data = OSX::NSData.alloc.initWithBytes_length(host_pem, host_pem.size)
    device_data = OSX::NSData.alloc.initWithBytes_length(device_pem, device_pem.size)
  
  
    p x.pair(root_data, device_data, host_data, "29942954112249261629052abcdef40")
  end
end

def activate
  devs = AppleDevice.available_devices
  
  if devs[0].kind_of?(NormalMode)
    x = devs[0].open
    
    #x.deactivate
    
    #abort
    
    uid = x.get_value("UniqueDeviceID")
    imei = x.get_value("InternationalMobileEquipmentIdentity")
    iccid = x.get_value("IntegratedCircuitCardIdentity")
    activation_info = x.get_value("ActivationInfo")
    serial_number = x.get_value("SerialNumber")
    imsi = x.get_value("InternationalMobileSubscriberIdentity")
    
    #p uid
    #p imei
    #p iccid
    #puts activation_info.to_plist
    #p serial_number
        
    #abort
    
    url = URI("https://albert.apple.com/WebObjects/ALActivation.woa/wa/iPhoneRegistration")

    http = Net::HTTP.new(url.host, url.port)

    http.use_ssl = true

    http.start do |h|
      req = Net::HTTP::Post.new(url.path, "User-Agent" => "iTunes/7.7 (Macintosh; U; Intel Mac OS X 10.5.4)")

      req.form_data = {
        "AppleSerialNumber" => serial_number,
        "IMSI" => imsi,
        "InStoreActivation" => "false",
        "machineName" => "macos",
        "activation-info" => activation_info.to_plist,
        "ICCID" => iccid,
        "IMEI" => imei
      }

      #puts req.body

      result = h.request(req)

      puts result.body

      document = REXML::Document.new(result.body)
      
      activation = OSX.load_plist(document.root.elements[1].elements[1].to_s)
      
      pp activation
      
      activation_record = activation["iphone-activation"]["activation-record"]
      
      x.activate(activation_record)
    end
  
  end
end



__END__

# Here's other random stuff

base_path = "/Users/pumpkin/Desktop/iPhone1,2_2.0_5A347_Restore"

if devs[0].kind_of?(RecoveryV2Mode)
  x = devs[0].open

  p "sending apple logo"
  
  File.open("#{base_path}/Firmware/all_flash/all_flash.n82ap.production/applelogo.s5l8900x.img3") do |f|
    x.send_file(f)
  end

  x.send_command("setpicture 1")
  x.send_command("bgcolor 0 255 0")

  p "sending device tree"
  File.open("#{base_path}/Firmware/all_flash/all_flash.n82ap.production/DeviceTree.n82ap.img3") do |f|
    x.send_file(f)
  end

  x.send_command("devicetree")

  pp x.recv_buffer.split("\x00")

  p "sending ramdisk"
  File.open("#{base_path}/018-3783-2.dmg") do |f|
    x.send_file(f)
  end

  x.send_command("ramdisk")

  pp x.recv_buffer.split("\x00")

  p "sending kernel"
  File.open("#{base_path}/kernelcache.release.s5l8900x") do |f|
    x.send_file(f)
  end

  p "booting"
  x.send_command("bootx")

  pp x.recv_buffer.split("\x00")

  p "sleeping"
  #sleep(60)
  
end

devs = AppleDevice.available_devices

p devs[0]

if devs[0].kind_of?(NormalMode)
  restore = devs[0].open
  
  progress_callback = proc do |operation, progress|
    steps = {
      28 => "Waiting for NAND",
      11 => "Creating partition map",
      12 => "Creating filesystem",
      13 => "Restoring image",
      14 => "Verifying restore",
      15 => "Checking filesystems",
      16 => "Mounting filesystems",
      29 => "Fixing up /var",
      29 => "Unmounting filesystems",
      25 => "Modifying persistent boot-args",
      35 => "Loading NOR data to flash",
      18 => "Flashing NOR",
      19 => "Updating baseband",
      20 => "Finalizing NAND epoch update",
      32 => "Waiting for Device..."
    }
    
    puts "#{steps[operation]} (#{operation}) with progress #{progress}"
  end
  
  p "starting restore"
  restore.start_restore(progress_callback) do |data_type|
    puts "DataRequest callback"
    
    if data_type == "SystemImageData"
      puts "Got request for system image data"
      
      Thread.new do
        puts "Started ASR thread" 
        File.open("#{base_path}/018-3782-2.dmg") do |f|
          asr = ASRService.new(12345, f)
          asr.start
        end
      end
      
      nil
    elsif data_type == "NORData"
      puts "Got request for NOR data"
      
      llb_data = File.read("#{base_path}/Firmware/all_flash/all_flash.n82ap.production/LLB.n82ap.RELEASE.img3")

      other_nor_data = Dir["#{base_path}/Firmware/all_flash/all_flash.n82ap.production/*.img3"].reject{|x| x =~ /^LLB/}.map do |path|
        buffer = File.read(path)
        OSX::NSData.alloc.initWithBytes_length(buffer, buffer.size)
      end

      response = { "LlbImageData" => OSX::NSData.alloc.initWithBytes_length(llb_data, llb_data.size),
                   "NorImageData" => other_nor_data }
      
      response
    end    
  end
  
  puts "Rebooting"
  
  
end

__END__