import re

import pymem.memory
import pymem.ressources.kernel32
import pymem.ressources.structure


def scan_pattern_page(handle, address, pattern):
    """Search a byte pattern given a memory location.
    Will query memory location information and search over until it reaches the
    length of the memory page. If nothing is found the function returns the
    next page location.

    Parameters
    ----------
    handle: HANDLE
        Handle to an open object
    address: int
        An address to search from
    pattern: bytes
        A regex byte pattern to search for

    Returns
    -------
    tuple
        next_region, found address

        found address may be None if one was not found or we didn't have permission to scan
        the region
    """
    mbi = pymem.memory.virtual_query(handle, address)
    next_region = mbi.BaseAddress + mbi.RegionSize
    allowed_protections = [
        pymem.ressources.structure.MEMORY_PROTECTION.PAGE_EXECUTE_READ,
        pymem.ressources.structure.MEMORY_PROTECTION.PAGE_EXECUTE_READWRITE,
        pymem.ressources.structure.MEMORY_PROTECTION.PAGE_READWRITE,
        pymem.ressources.structure.MEMORY_PROTECTION.PAGE_READONLY,
    ]
    if mbi.state != pymem.ressources.structure.MEMORY_STATE.MEM_COMMIT or mbi.protect not in allowed_protections:
        return next_region, None

    page_bytes = pymem.memory.read_bytes(handle, address, mbi.RegionSize)

    found = None

    match = re.search(pattern, page_bytes, re.DOTALL)

    if match:
        found = address + match.span()[0]

    return next_region, found


def pattern_scan_module(handle, module, pattern):
    """Given a handle over an opened process and a module will scan memory after
    a byte pattern and return its corresponding memory address.

    Parameters
    ----------
    handle: HANDLE
        Handle to an open object
    module: MODULEINFO
        An instance of a given module
    pattern: bytes
        A regex byte pattern to search for

    Returns
    -------
    Optional[int]
        Memory address of given pattern, or None if one was not found

    Examples
    --------
    >>> pm = pymem.Pymem("Notepad.exe")
    # Here the "." means that the byte can be any byte; a "wildcard"
    # also note that this pattern may be outdated
    >>> character_count_pattern = b".\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00" \
    ...                           b"\x00\x00\x00\x00\x00\x00..\x00\x00..\x00\x00\x64\x04"
    >>> module = pymem.process.module_from_name(pm.process_handle, "Notepad.exe")
    >>> character_count_address = pymem.pattern.pattern_scan_module(pm.process_handle, module, character_count_pattern)
    """
    base_address = module.lpBaseOfDll
    max_address = module.lpBaseOfDll + module.SizeOfImage
    page_address = base_address

    found = None
    while page_address < max_address:
        next_page, found = scan_pattern_page(handle, page_address, pattern)
        if found:
            break
        page_address = next_page

    return found

@@ -131,32 +131,58 @@
 address, pattern
+, *, return_multiple=False
 ):%0A    %22%22%22Search
@@ -589,32 +589,138 @@
 rn to search for
+%0A    return_multiple: bool%0A        If multiple results should be returned instead of stopping on the first
 %0A%0A    Returns%0A  
@@ -885,16 +885,147 @@
 e region
+%0A%0A        if return_multiple is True found address will instead be a list of found addresses%0A        or an empty list if no results
 %0A    %22%22%22
@@ -1673,16 +1673,48 @@
 nSize)%0A%0A
+    if not return_multiple:%0A    
     foun
@@ -1718,25 +1718,28 @@
 ound = None%0A
-%0A
+    
     match = 
@@ -1781,16 +1781,20 @@
 OTALL)%0A%0A
+    
     if m
@@ -1799,32 +1799,36 @@
  match:%0A        
+ 
+   
 found = address 
@@ -1846,16 +1846,207 @@
 n()%5B0%5D%0A%0A
+    else:%0A        found = %5B%5D%0A%0A        for match in re.finditer(pattern, page_bytes, re.DOTALL):%0A            found_address = address + match.span()%5B0%5D%0A            found.append(found_address)%0A%0A
     retu
@@ -2116,16 +2116,42 @@
  pattern
+, *, return_multiple=False
 ):%0A    %22
@@ -2495,16 +2495,122 @@
 arch for
+%0A    return_multiple: bool%0A        If multiple results should be returned instead of stopping on the first
 %0A%0A    Re
@@ -2631,16 +2631,22 @@
 ---%0A    
+Union%5B
 Optional
@@ -2650,16 +2650,23 @@
 nal%5Bint%5D
+, list%5D
 %0A       
@@ -2727,16 +2727,80 @@
 ot found
+%0A        or a list of found addresses in return_multiple is True
 %0A%0A    Ex
@@ -3496,24 +3496,56 @@
 se_address%0A%0A
+    if not return_multiple:%0A    
     found = 
@@ -3545,24 +3545,28 @@
 ound = None%0A
+    
     while pa
@@ -3591,16 +3591,20 @@
 ddress:%0A
+    
         
@@ -3671,24 +3671,26 @@
 attern)%0A
+%0A 
         
 if found
@@ -3681,16 +3681,19 @@
         
+   
 if found
@@ -3710,14 +3710,301 @@
     
+    
 break%0A
+%0A            page_address = next_page%0A%0A    else:%0A        found = %5B%5D%0A        while page_address %3C max_address:%0A            next_page, new_found = scan_pattern_page(handle, page_address, pattern, return_multiple=True)%0A%0A            if new_found:%0A                found += new_found%0A%0A    
     

add return_multiple keyword for pattern functions (#59)