#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals

import argparse
import codecs
import os
import struct

import gdb
import pwndbg.color.memory as M
import pwndbg.commands
import pwndbg.enhance
import pwndbg.search
import pwndbg.vmmap


def print_search(value, *a, **kw):
    hits = set()

    for address in pwndbg.search.search(value, *a, **kw):
        if not address:
            continue

        if address in hits:
            continue

        hits.add(address)

        vmmap = pwndbg.vmmap.find(address)
        if vmmap:
            region = os.path.basename(vmmap.objfile)
        else:
            region = '[mapped]'

        region = region.ljust(15)

        region = M.get(address, region)
        addr = M.get(address)
        display = pwndbg.enhance.enhance(address)
        print(region,addr,display)

parser = argparse.ArgumentParser(description='''
Search memory for byte sequences, strings, pointers, and integer values
''')
parser.add_argument('-t', '--type', choices=['byte','short','dword','qword','pointer','string','bytes'],
                    help='Size of search target', default='bytes', type=str)
parser.add_argument('-1', '--byte', dest='type', action='store_const', const='byte',
                    help='Search for a 1-byte integer')
parser.add_argument('-2', '--word', dest='type', action='store_const', const='word',
                    help='Search for a 2-byte integer')
parser.add_argument('-4', '--dword', dest='type', action='store_const', const='dword',
                    help='Search for a 4-byte integer')
parser.add_argument('-8', '--qword', dest='type', action='store_const', const='qword',
                    help='Search for an 8-byte integer')
parser.add_argument('-p', '--pointer', dest='type', action='store_const', const='pointer',
                    help='Search for a pointer-width integer')
parser.add_argument('-x', '--hex', action='store_true',
                    help='Target is a hex-encoded (for bytes/strings)')
parser.add_argument('-s', '--string', action='store_true',
                    help='Target is a raw string')
parser.add_argument('-e', '--executable', action='store_true',
                    help='Search executable segments only')
parser.add_argument('-w', '--writable', action='store_true',
                    help='Search writable segments only')
parser.add_argument('value', type=str,
                    help='Value to search for')
parser.add_argument('mapping', type=str, nargs='?', default=None,
                    help='Mapping to search [e.g. libc]')

@pwndbg.commands.ArgparsedCommand(parser)
@pwndbg.commands.OnlyWhenRunning
def search(type, hex, string, executable, writable, value, mapping):
    # Adjust pointer sizes to the local architecture
    if type == 'pointer':
        type = {
            4: 'dword',
            8: 'qword'
        }[pwndbg.arch.ptrsize]

    if hex:
        value = codecs.decode(value, 'hex')

    # Convert to an integer if needed, and pack to bytes
    if type not in ('string', 'bytes'):
        value = pwndbg.commands.fix_int(value)
        fmt = {
            'little': '<',
            'big': '>'
        }[pwndbg.arch.endian] + {
            'byte': 'B',
            'short': 'H',
            'dword': 'L',
            'qword': 'Q'
        }[type]

        value = struct.pack(fmt, value)
    
    # Null-terminate strings
    elif type == 'string':
        value += b'\x00'

    # Perform the search
    print_search(value, mapping=mapping, executable=executable, writable=writable)

@@ -303,16 +303,30 @@
 .vmmap%0A%0A
+saved = set()%0A
 %0Adef pri
@@ -338,175 +338,143 @@
 arch
-(value, *a, **kw):%0A    hits = set()%0A%0A    for address in pwndbg.search.search(value, *a, **kw):%0A        if not address:%0A            continue%0A
+_hit(address):%0A    %22%22%22Prints out a single search hit.%0A%0A    Arguments:%0A        address(int): Address to print
 %0A    
+%22%22%22%0A
     if 
+not 
 address
- in hits
 :%0A  
@@ -483,53 +483,16 @@
     
-    continue%0A%0A        hits.add(address)%0A%0A    
+return%0A%0A
     
@@ -522,28 +522,24 @@
 nd(address)%0A
-    
     if vmmap
@@ -540,36 +540,32 @@
  vmmap:%0A        
-  
-  
 region = os.path
@@ -589,20 +589,16 @@
 bjfile)%0A
-    
     else
@@ -603,28 +603,24 @@
 se:%0A        
-  
-  
 region = '%5Bm
@@ -628,20 +628,16 @@
 pped%5D'%0A%0A
-    
     regi
@@ -659,28 +659,24 @@
 st(15)%0A%0A    
- 
-   
 region = M.g
@@ -695,20 +695,16 @@
 region)%0A
-    
     addr
@@ -725,20 +725,16 @@
 ss)%0A    
- 
-   
 display 
@@ -767,20 +767,16 @@
 ddress)%0A
-    
     prin
@@ -2520,16 +2520,268 @@
  libc%5D')
+%0Aparser.add_argument('--save', action='store_true',%0A                    help='Save results for --resume')%0Aparser.add_argument('-n', '--next', action='store_true',%0A                    help='Search only locations returned by previous search with --save')
 %0A%0A@pwndb
@@ -2915,16 +2915,28 @@
  mapping
+, save, next
 ):%0A    #
@@ -3565,20 +3565,16 @@
  value)%0A
-    
 %0A    # N
@@ -3659,105 +3659,492 @@
  # P
-erform the search%0A    print_search(value, mapping=mapping, executable=executable, writable=writable
+rep the saved set if necessary%0A    global saved%0A    if save:%0A        saved = set()%0A    %0A    # Perform the search%0A    for address in pwndbg.search.search(value,%0A                                        mapping=mapping,%0A                                        executable=executable,%0A                                        writable=writable):%0A%0A        if next and address not in saved:%0A            continue%0A%0A        if save:%0A            saved.add(address)%0A%0A        print_search_hit(address
 )%0A

Add CheatEngine search with --save and --next