#-*- coding: UTF-8 -*-
# Coded by Orange@chroot.org

import wx
import os
import sys
import time
import runpy
import socket
import shutil
import gettext
import threading

import wx.lib.agw.pyprogress as PP
from BeautifulSoup import BeautifulSoup

# Set script PATH
sys.path.append("./usbmuxd")
sys.path.append("./scripts")

from util.ramdiskclient import RamdiskToolClient
from keystore.keybag import Keybag

# set i18n
_ = wx.GetTranslation

# Hook debug window
class LogWindow(wx.PyOnDemandOutputWindow):
    def __init__(self):
        wx.PyOnDemandOutputWindow.__init__(self, title = "Message Window")

    def write(self, text):
        if self.frame is None:

            # Bug Fix with Init self.frame fast... XD
            self.frame = "____"
            if not wx.Thread_IsMain():
                wx.CallAfter(self.CreateOutputWindow, text)
            else:
                self.CreateOutputWindow(text)
        else:
            if not wx.Thread_IsMain():
                wx.CallAfter(self.__write, text)
            else:
                self.__write(text)
                
    def __write(self, text):
        try:
            self.text.AppendText(text)
        except AttributeError:
            pass

        for item in self.triggers:
            if item in text:
                self.frame.Raise()
                break

# hook print statement
class HookStdOut(object):
    def __init__(self, *argv, **argd):
        object.__init__(self)

    def write(self, s):
        try:
            sys.__stdout__.write(s)
        except:
            pass

class MyApp(wx.App):
    outputWindowClass = LogWindow
    def __init__(self, redirect=True):
        wx.App.__init__(self, redirect)

class JB_Dialog(wx.Dialog):
    def __init__(self, parent, ID, title, size=wx.DefaultSize, pos=wx.DefaultPosition, style=wx.DEFAULT_DIALOG_STYLE,):
        pre = wx.PreDialog()
        pre.Create(parent, ID, title, pos, size, style)
        self.PostCreate(pre)

        # set description str
        label = wx.StaticText(self, -1, _("\nPlease select your decide type."))
        sizer = wx.BoxSizer(wx.VERTICAL)
        sizer.Add(label, 0, wx.ALIGN_CENTRE|wx.ALL, 5)

        # set select bar
        self.l = wx.Choice(self,
                           -1,
                           (100, 50),
                           choices=["Iphone 4 (GSM)", "Iphone 4 (CDMA)", "Iphone 3GS ", "Ipod Touch 3G", "Ipod Touch 4G", "Ipad 1"]
                          )
        self.l.SetSelection(0)
        sizer.Add(self.l, 2, wx.EXPAND|wx.ALL, 10)
        box = wx.BoxSizer(wx.HORIZONTAL)
        btnsizer = wx.StdDialogButtonSizer()

        # set OK button
        btn = wx.Button(self, wx.ID_OK)
        btn.SetDefault()
        btnsizer.AddButton(btn)

        # set CANCEL button
        btn = wx.Button(self, wx.ID_CANCEL)
        btnsizer.AddButton(btn)
        btnsizer.Realize()

        sizer.Add(btnsizer, 0, wx.ALIGN_CENTER_VERTICAL|wx.ALL, 5)
        self.SetSizer(sizer)
        sizer.Fit(self)

class BruteForce_Dialog(wx.Dialog):
    def __init__(self, parent, ID, title, size=wx.DefaultSize, pos=wx.DefaultPosition, style=wx.DEFAULT_DIALOG_STYLE,):
        pre = wx.PreDialog()
        pre.Create(parent, ID, title, pos, size, style)
        self.PostCreate(pre)

        # set description str
        sizer = wx.BoxSizer(wx.VERTICAL)

        # set select bar
        self.l = wx.Choice(self,
                           -1,
                           (100, 50),
                           choices=[_("The password is empty."), _("I know the password."), _("I do not know the password.")]
                          )
        self.l.SetSelection(0)
        sizer.Add(self.l, 2, wx.EXPAND|wx.ALL, 10)
        box = wx.BoxSizer(wx.HORIZONTAL)
        btnsizer = wx.StdDialogButtonSizer()

        # set OK button
        btn = wx.Button(self, wx.ID_OK)
        btn.SetDefault()
        btnsizer.AddButton(btn)

        # set CANCEL button
        btn = wx.Button(self, wx.ID_CANCEL)
        btnsizer.AddButton(btn)
        btnsizer.Realize()

        sizer.Add(btnsizer, 0, wx.ALIGN_CENTER_VERTICAL|wx.ALL, 5)
        self.SetSizer(sizer)
        sizer.Fit(self)


class LangChoice(wx.Dialog):
    def __init__(self, parent, ID, title, size=wx.DefaultSize, pos=wx.DefaultPosition, style=wx.DEFAULT_DIALOG_STYLE,):
        pre = wx.PreDialog()
        pre.Create(parent, ID, title, pos, size, style)
        self.PostCreate(pre)
        sizer = wx.BoxSizer(wx.VERTICAL)

        # set select bar
        self.l = wx.Choice(self,
                           -1,
                           (100, 50),
                           choices=["Chinese", "English", "Japanese"]
                          )
        self.l.SetSelection(0)
        sizer.Add(self.l, 2, wx.EXPAND|wx.ALL, 10)
        box = wx.BoxSizer(wx.HORIZONTAL)
        btnsizer = wx.StdDialogButtonSizer()

        # set OK button
        btn = wx.Button(self, wx.ID_OK)
        btn.SetDefault()
        btnsizer.AddButton(btn)

        # set CANCEL button
        btn = wx.Button(self, wx.ID_CANCEL)
        btnsizer.AddButton(btn)
        btnsizer.Realize()

        sizer.Add(btnsizer, 0, wx.ALIGN_CENTER_VERTICAL|wx.ALL, 5)
        self.SetSizer(sizer)
        sizer.Fit(self)

class main(wx.Frame):
    def __init__(self, parent, id, title, size):
        wx.Frame.__init__(self, parent, id, title, size = size)
        self._main()

    def CreateElement(self):
        panel = wx.Panel(self, -1, style=wx.RAISED_BORDER)
        self.option1 = wx.Button(panel, -1, _("Change language"), style=wx.NO_BORDER)
        self.option2 = wx.Button(panel, -1, _("Using SSH over USB, Setup Usbmux. (Optional)"), style=wx.NO_BORDER)
        self.option3 = wx.Button(panel, -1, _("Using JB Exploit to Load Forensic Ramdisk."), style=wx.NO_BORDER)
        self.option4 = wx.Button(panel, -1, _("Brute Force the Passcode."), style=wx.NO_BORDER)
        self.option5 = wx.Button(panel, -1, _("Export the Important Password Report from Keychain."), style=wx.NO_BORDER)
        self.option6 = wx.Button(panel, -1, _("Mirror IOS Data Partition."), style=wx.NO_BORDER)
        self.option7 = wx.Button(panel, -1, _("Decrypt IOS Data Pratition Image."), style=wx.NO_BORDER)
        self.option8 = wx.Button(panel, -1, _("Recover Deleted Files from Decrypted Image."), style=wx.NO_BORDER)
        self.option9 = wx.Button(panel, -1, _("Deeply Revocer Deleted Files from Decrypted Image."), style=wx.NO_BORDER)

        box = wx.BoxSizer(wx.VERTICAL)

        # dirty way to eval button
        for i in range(1, 10):
            eval("box.Add(self.option%d, 5, wx.EXPAND | wx.ALL, 4)" % i)

        panel.SetSizer(box)

    def OnShow(self, evt):
        value = wx.GetTextFromUser(_("\nPlease enter case number (12 digits)."), _("Message!"), "")
        if value:
            try:
                value = int(value)
            except:
                wx.MessageBox("Case number must be digit.", _("Error!"), style=wx.ICON_HAND)
                value = None

        if value:
            # do_lookup(value)
            pass

    def ChgLang(self, evt):
        dlg = LangChoice(self, -1, _("Change Language"), size=(350, 200))
        dlg.CenterOnScreen()
        val = dlg.ShowModal()
        if val == wx.ID_OK:
            choice = dlg.l.CurrentSelection
            if choice == 0:
                mylocale.AddCatalog('cht')
            elif choice == 1:
                mylocale.AddCatalog('en')
            elif choice == 2:
                mylocale.AddCatalog('jp')


        dlg.Destroy()

    def JB(self, evt):
        dlg = JB_Dialog(self, -1, _("Using JB Exploit to Load Forensic Ramdisk"), size=(350, 200))
        dlg.CenterOnScreen()
        val = dlg.ShowModal()

        if val == wx.ID_OK:
            choice = dlg.l.CurrentSelection
            if choice == 0:     # Iphone 4 (GSM)
                i = "iPhone3,1_5.0_9A334_Restore.ipsw"
                r = "iPhone3,1_5.0_9A334_Restore.dmg"
                k = "iPhone3,1_5.0_9A334_Restore.patched"
            elif choice == 1:   # Iphone 4 (CDMA)
                i = "iPhone3,3_5.0_9A334_Restore.ipsw"
                r = "iPhone3,3_5.0_9A334_Restore.dmg"
                k = "iPhone3,3_5.0_9A334_Restore.patched"
            elif choice == 2:   # Iphone 3GS
                i = "iPhone2,1_5.0_9A334_Restore.ipsw"
                r = "iPhone2,1_5.0_9A334_Restore.dmg"
                k = "iPhone2,1_5.0_9A334_Restore.patched"
            elif choice == 3:   # Ipod Touch 3G
                i = "iPod3,1_5.0_9A334_Restore.ipsw"
                r = "iPod3,1_5.0_9A334_Restore.dmg"
                k = "iPod3,1_5.0_9A334_Restore.patched"
            elif choice == 4:   # Ipod Touch 4G
                i = "iPod4,1_5.0_9A334_Restore.ipsw"
                r = "iPod4,1_5.0_9A334_Restore.dmg"
                k = "iPod4,1_5.0_9A334_Restore.patched"
            elif choice == 5:   # Ipad 1
                i = "iPad1,1_5.0_9A334_Restore.ipsw"
                r = "iPad1,1_5.0_9A334_Restore.dmg"
                k = "iPad1,1_5.0_9A334_Restore.patched"
            
            # run redsn0w in order to JailBreak
            wx.MessageBox(_("Please click CANCEL button to reboot after JailBreak finish."), _("Message!"))
            os.system(".\\tools\\redsn0w.exe -i .\\dmg\\%s -r .\\dmg\\%s -k .\\dmg\\%s" % (i, r, k))


        dlg.Destroy()

    def MountUSB(self, evt):
        if not CheckPort(2222):
            threading.Thread(target=self._MountUSB).start()

    def _MountUSB(self):
        # Run script
        sys.argv = []
        sys.argv.append( "" )
        sys.argv.append( "-t" )
        sys.argv.append( "22:2222" )
        sys.argv.append( "1999:1999" )

        runpy.run_module("tcprelay", run_name="__main__", alter_sys=True, init_globals={"raw_input": msg})

    def BruteForce(self, evt):
        # Check tcp relay work
        if not CheckPort(2222):
            self.MountUSB(None)

        threading.Thread(target=self._BruteForce).start()
        
        """
        dlg = BruteForce_Dialog(self, -1, _("Message"), size=(350, 200))
        dlg.CenterOnScreen()
        val = dlg.ShowModal()

        if val == wx.ID_OK:
            choice = dlg.l.CurrentSelection
            # Know password
            if choice == 0:
                threading.Thread(target=self._BruteForce).start()
            else:
                threading.Thread(target=self._BruteForce).start()
        dlg.Destroy()
        """

    def _BruteForce(self):
        # Run script
        runpy.run_module("demo_bruteforce", run_name="__main__", alter_sys=True)

        filename = "%s.plist" % GetDi("dataVolumeUUID")
        if os.path.exists(filename):
            with open(filename, "a+") as fp:
                flag = False
                for i in fp.readlines():
                    if flag:
                        pw = i
                        flag = False
                        break

                    if i.find("<key>passcode</key>") != -1:
                        flag = True
    
                if pw:
                    pw = BeautifulSoup(pw)
                    pw = pw.findAll("string")[0].string
                    if pw == None:
                        wx.MessageBox(_("There is no password") , _("Message!"))
                    else:
                        wx.MessageBox(_("The password is %s") % pw, ("Message!"))
        else:
            wx.MessageBox(_("Can not find password file"), _("Error!"))
        
        # Back to work direcytory
        os.chdir("../")

    def KeychainDeCrypt(self, evt):
        # Check tcp relay work
        # QQ strange bug QQ
        if not CheckPort(2222):
            self.MountUSB(None)

        threading.Thread(target=self._KeychainDeCrypt).start()

    def _KeychainDeCrypt(self):
        # Check plist and keychain
        if os.path.exists( "%s/%s.plist" % (GetDi("udid"), GetDi("dataVolumeUUID") ) ) and os.path.exists( "%s/keychain-2.db" % GetDi("udid") ):
            # Run script
            sys.argv = []
            sys.argv.append( "" )
            sys.argv.append( "-d")
            sys.argv.append( ".\\%s\\keychain-2.db" % GetDi("udid") )
            sys.argv.append( ".\\%s\\%s.plist" % (GetDi("udid"), GetDi("dataVolumeUUID")) )
            runpy.run_module("keychain_tool", run_name="__main__", alter_sys=True, init_globals={"raw_input": msg})

            wx.MessageBox(_("The password report has been saved."), _("Message!"))
        else:
            wx.MessageBox(_("Please brute force the passcode first."), _("Message!"))

    def MirrorIOS(self, evt):
        # Check tcp relay work
        if not CheckPort(2222):
            self.MountUSB(None)


        wx.MessageBox(_("Please switch to the terminal window, and enter password <alpine>."), _("Notice"))

        tmp = {"flag": True}
        threading.Thread(target=self._MirrorIOS, args=(tmp,)).start()
        dlg = PP.PyProgress(None, -1, "Dialog", "Progress")
        while tmp["flag"]:
            wx.MilliSleep(250)
            mb = os.path.getsize( "%s/encrypted.img" % GetDi("udid") ) / (1000*1000)
            dlg.Refresh()
            dlg.UpdatePulse("%d MB" % mb)
        dlg.Destroy()

        if os.path.getsize( "%s/encrypted.img" % GetDi("udid") ) == 0:
            wx.MessageBox(_("Failed, May the password is incorrect."), _("Message"))
        else:
            wx.MessageBox(_("Mirroring finish, the file name is encrypted.img."), _("Message"))

    def _MirrorIOS(self, arg):
        class __(object):
            def __init__(self, *argv, **argd):
                object.__init__(self)

            def write(self, s):
                pass
        sys.stdout = __()
        cmd = "test -e /dev/rdisk0s2s1 && a='rdisk0s2s1';test -e /dev/rdisk0s1s2 && a='rdisk0s1s2';dd if=/dev/$a bs=8192;"
        os.system(".\\ssh\ssh.exe -p 2222 root@localhost \"%s\" | .\\ssh\\dd.exe of=%s/encrypted.img" % (cmd, GetDi("udid")) )
        sys.stdout = HookStdOut()
        arg["flag"] = False

    def MirrorDecrypt(self, evt):
        # Check tcp relay work
        if not CheckPort(2222):
            self.MountUSB(None)

        os.chdir(GetDi("udid"))

        # Check image.img exists
        if not os.path.exists("encrypted.img") and not os.path.exists("decrypted.img"):
            global _
            wx.MessageBox(_("Please Do Mirror Action First."), _("Message"))
            os.chdir("../")
            return None

        if not os.path.exists("backup_encrypted.img" ):
            def __(arg):
                shutil.copy("encrypted.img", "backup_encrypted.img")
                arg["flag"] = False

            # pass by reference so using a list
            tmp = {"flag": True}
            threading.Thread(target=__, args=(tmp,)).start()
            dlg = PP.PyProgress(None, -1, _("Dialog"), _("Please Watting Backup the Image File."))

            while tmp["flag"]:
                wx.MilliSleep(250)
                progress = (os.path.getsize("backup_encrypted.img")*100.0) / os.path.getsize("encrypted.img")
                dlg.Refresh()
                dlg.UpdatePulse("%d %%" % progress)
            dlg.Destroy()

        os.chdir("../")
        threading.Thread(target=self._MirrorDecrypt).start()

    def _MirrorDecrypt(self):
        if os.path.exists("%s/encrypted.img" % GetDi("udid")):
            os.rename( "%s/encrypted.img" % GetDi("udid"), "%s/decrypted.img" % GetDi("udid") )

        # Run script
        sys.argv = []
        sys.argv.append( "" )
        sys.argv.append( "%s/decrypted.img" % GetDi("udid") )
        runpy.run_module("emf_decrypter", run_name="__main__", alter_sys=True, init_globals={"raw_input": msg})

        wx.MessageBox(_("Finish IOS mirror decryption, original image file name is encrypted.img, decrypted file name is decryptde.img."), _("Message"))

    def MirrorUnDelete(self, evt):
        # Check tcp relay work
        if not CheckPort(2222):
            self.MountUSB(None)

        # Check image.img exists
        if not os.path.exists( "%s/decrypted.img" % GetDi("udid") ):
            global _
            wx.MessageBox(_("Please Do Mirror Action First."), _("Message"))
            return None

        threading.Thread(target=self._MirrorUnDelete).start()

    def _MirrorUnDelete(self):
        # Run script
        sys.argv = []
        sys.argv.append( "" )
        sys.argv.append( "%s/decrypted.img" % GetDi("udid") )
        runpy.run_module("emf_undelete", run_name="__main__", alter_sys=True, init_globals={"raw_input": msg})

        wx.MessageBox(_("Recover files are under directory %s/%s ." % (os.getcwd(), GetDi("udid"))), _("Message"))

    def MirrorUnDeleteDeep(self, evt):
        # Check tcp relay work
        if not CheckPort(2222):
            self.MountUSB(None)

        # Check image.img exists
        if not os.path.exists( "%s/decrypted.img" % GetDi("udid") ):
            global _
            wx.MessageBox(_("Please Do Mirror Action First."), _("Message"))
            return None
            
        threading.Thread(target=self._MirrorUnDeleteDeep).start()

    def _MirrorUnDeleteDeep(self):
        # Run script
        sys.argv = []
        sys.argv.append( "" )
        sys.argv.append( "%s/decrypted.img" % GetDi("udid") )
        runpy.run_module("Undelete_deep", run_name="__main__", alter_sys=True)

    def _main(self):
        # Create Buttons & Texts
        self.CreateElement()

        # Binding events
        self.option1.Bind(wx.EVT_BUTTON, self.ChgLang)
        self.option2.Bind(wx.EVT_BUTTON, self.MountUSB)
        self.option3.Bind(wx.EVT_BUTTON, self.JB)
        self.option4.Bind(wx.EVT_BUTTON, self.BruteForce)
        self.option5.Bind(wx.EVT_BUTTON, self.KeychainDeCrypt)
        self.option6.Bind(wx.EVT_BUTTON, self.MirrorIOS)
        self.option7.Bind(wx.EVT_BUTTON, self.MirrorDecrypt)
        self.option8.Bind(wx.EVT_BUTTON, self.MirrorUnDelete)
        self.option9.Bind(wx.EVT_BUTTON, self.MirrorUnDeleteDeep)
        #self.Bind(wx.EVT_SHOW, self.OnShow)

        self.Centre()
        self.Show()


def CheckPort(port=2222):
    ret = False

    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.bind( ("127.0.0.1", port) )
        s.close()
    except Exception, e:
        ret = True

    return ret
    
def GetDi(key):
    client = RamdiskToolClient()
    di = client.getDeviceInfos()

    #udid or dataVolumeUUID
    return di[key]

# Hook print statement
sys.stdout = HookStdOut()

# Hook raw_input()
def msg():
    wx.MessageBox(_("Plese enter OK to continue."), _("Message"))
raw_input = msg

if __name__ == "__main__":
    app = MyApp(redirect=True)

    # do i18n
    mylocale = wx.Locale()
    mylocale.AddCatalogLookupPathPrefix("./lang/")
    mylocale.AddCatalog('cht')

    frame = main(None, -1, _("IOS  Forensic Tools for Wiindows"), (400, 350))
    app.MainLoop()