# Author: Nicolas VERDIER
# This file is part of memorpy.
#
# memorpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# memorpy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with memorpy.  If not, see <http://www.gnu.org/licenses/>.

import copy
import time
import struct

from Address import Address


class Locator:
    """
    take a memoryworker and a type to search then you can feed the locator
    with values and it will reduce the addresses possibilities
    """

    def __init__(self, mw, type="unknown", start=None, end=None):
        self.mw = mw
        self.type = type
        self.last_iteration = {}
        self.last_value = None
        self.start = start
        self.end = end

    def find(self, value, erase_last=True):
        return self.feed(value, erase_last)

    def feed(self, value, erase_last=True):
        self.last_value = value
        new_iter = copy.copy(self.last_iteration)
        if self.type == "unknown":
            all_types = [
                "uint",
                "int",
                "long",
                "ulong",
                "float",
                "double",
                "short",
                "ushort",
            ]
        else:
            all_types = [self.type]
        for type in all_types:
            if type not in new_iter:
                try:
                    new_iter[type] = [
                        Address(x, self.mw.process, type)
                        for x in self.mw.mem_search(
                            value, type, start_offset=self.start, end_offset=self.end
                        )
                    ]
                except struct.error:
                    new_iter[type] = []
            else:
                l = []
                for address in new_iter[type]:
                    try:
                        found = self.mw.process.read(address, type)
                        if int(found) == int(value):
                            l.append(Address(address, self.mw.process, type))
                    except Exception as e:
                        pass

                new_iter[type] = l

        if erase_last:
            del self.last_iteration
            self.last_iteration = new_iter
        return new_iter

    def get_addresses(self):
        return self.last_iteration

    def diff(self, erase_last=False):
        return self.get_modified_addr(erase_last)

    def get_modified_addr(self, erase_last=False):
        last = self.last_iteration
        new = self.feed(self.last_value, erase_last=erase_last)
        ret = {}
        for type, l in last.iteritems():
            typeset = set(new[type])
            for addr in l:
                if addr not in typeset:
                    if type not in ret:
                        ret[type] = []
                    ret[type].append(addr)

        return ret
