Working on data / fixing inconsistency in code Xref

__init__.py

@@ -8,8 +8,9 @@
 import xref
 import ida_import
 import idb
+import data
 
-all_submodules_name = ['elt', 'functions', 'xref', 'ida_import', 'idb']
+all_submodules_name = ['elt', 'functions', 'xref', 'ida_import', 'idb', 'data']
 
 
 def get_full_submodule_name(name):
@@ -21,12 +22,12 @@ def get_full_submodule_name(name):
 #TODO : You know what for problem with kernel32.dll
 
 #TODO data.py -> Bytes / Word / Dword / String / cstr
+# ?? : Undef Data et (Undef code : exist now) ?
 
 
 # helper : reload(MIDAP); MIDAP.reload(); g = MIDAP.functions.MFunctions(); f = MIDAP.fhere(); i = MIDAP.ihere()
 
 # TODO: structures ?
-# Entry point: list(Entries())
 
 
 #TODO : CLEAR ON RELOAD
@@ -65,10 +66,14 @@ def bhere():
 def fhere():
     "Typed here(): return current Function"
     return functions.IDAFunction.get_func(idc.here())
+
+def dhere():
+    "Typed here(): return current Data"
+    return data.Data.new_data_by_type(idc.here())
 
 
 
 
 fixup_late_import()
 
-self = idb.IDB()
+self = idb.current

data.py

@@ -1,2 +1,107 @@
-class Data(object):
-    pass
+import elt
+
+import idc
+
+# why not: automatic subclassing :)
+class Data(elt.IDANamedSizedElt):
+
+    # size == 0 doesn't make sens: need automatic subclassing :D
+    size = 0
+    match = (lambda *args : False)
+    # addr -> bool
+    # staticmethod
+    _get_value_func = None
+
+    def value(self):
+        """ 
+            [property] [get | set]
+            The value of the Data object.
+            DataByte.value = 0x42 will patch the byte at DataByte.addr
+        """
+        return self._get_value_func(self.addr)
+
+    def set_value(self, value, litte_endian=True):
+        # TODO: ask endianness to IDB
+        initial_value = value
+        bytes = []
+        for i in range(self.size):
+            bytes.append(value & 255)
+            value = value >> 8
+        if value != 0:
+            raise ValueError("value {0} is too big for {1}".format(hex(initial_value), self.__class__.__name__))
+        if not litte_endian:
+            bytes.reverse()
+        self.patch(bytes, False)
+
+    value = property(value, set_value, None)
+
+
+    def __init__(self, addr):
+        super(Data, self).__init__(addr, addr + self.size)
+
+    @property
+    def is_byte(self):
+        return idc.isByte(self.flags)
+
+    @property
+    def is_word(self):
+        return idc.isWord(self.flags)
+
+    @property
+    def is_dword(self):
+        return idc.isDwrd(self.flags)
+
+    @property
+    def is_qword(self):
+        return idc.isQwrd(self.flags)
+
+    def __IDA_repr__(self):
+        if self._get_value_func is None:
+            return ""
+        return hex(self.value)
+
+  #  @property
+  #  def next(self):
+  #      return Data.new_data_by_type(
+
+    @classmethod
+    def new_data_by_size(cls, addr, size):
+        for subcls in cls.__subclasses__():
+            if subcls.size == size:
+                return subcls(addr)
+        raise ValueError("Don't know how to handle size = {0}".format(size))
+
+    @classmethod    
+    #Name new_data ? (most general use)
+    def new_data_by_type(cls, addr):
+        data = Data(addr)
+        for subcls in cls.__subclasses__():
+            if subcls.match(data):
+                return subcls(addr)
+        raise ValueError("Don't know how to handle addr = {0}".format(hex(addr)))
+
+
+
+class ByteData(Data):
+    size = 1
+    _get_value_func = staticmethod(idc.Byte)
+    match = staticmethod(Data.is_byte.fget)
+
+class WordData(Data):
+    size = 2
+    _get_value_func = staticmethod(idc.Word)
+    match = staticmethod(Data.is_word.fget)
+
+class DwordData(Data):
+    size = 4
+    _get_value_func = staticmethod(idc.Dword)
+    match = staticmethod(Data.is_dword.fget)
+
+
+class DwordData(Data):
+    size = 8
+    _get_value_func = staticmethod(idc.Qword)
+    match = staticmethod(Data.is_qword.fget)
+
+
+

elt.py

@@ -1,7 +1,7 @@
 import idc
 import idautils
 
-late_import = ['xref']
+late_import = ['xref', 'data']
 
 
 class IDAElt(object):
@@ -13,6 +13,10 @@ def get_addr(self):
 
     addr = property(get_addr, None, None, 'ea of the object')
 
+    @classmethod
+    def get_all(cls):
+        return [cls(x) for x in idautils.Heads()]
+
     #property ?
     def goto(self):
         idc.Jump(self.addr)
@@ -21,21 +25,84 @@ def __int__(self):
         return self.addr
 
     def __repr__(self):
-        return "<{0}>".format(self.__IDA_repr__() + " <at {0}>".format(hex(self.addr)))
+
+        return "<{cls} {ida_repr} <at {addr}>>".format(
+                cls=self.__class__.__name__,
+                ida_repr=self.__IDA_repr__(),
+                addr=hex(self.addr))
 
     def __IDA_repr__(self):
-        return self.__class__.__name__
+        return ""
 
     @property    
     def xfrom(self):
+        """ List of all XrefsFrom the element """
         return [xref.Xref(x) for x in idautils.XrefsFrom(self.addr, False)]
 
     @property   
     def xto(self):
+        """ List of all XrefsTo the element """
         return [xref.Xref(x) for x in idautils.XrefsTo(self.addr, False)]
 
-    #Create data xref ?
-
+    @property
+    def flags(self):
+        return idc.GetFlags(self.addr)
+
+    # do LineA and LineB ? for comments ?
+
+    # Flags
+
+    @property
+    def is_code(self):
+        return idc.isCode(self.flags)
+
+    @property
+    def is_data(self):
+        return idc.isData(self.flags)
+
+    @property    
+    def is_unknow(self):
+        return idc.isUnknown(self.flags)
+
+    @property
+    def is_head(self):
+        return idc.isHead(self.flags)
+
+    @property
+    def is_tail(self):
+        return idc.isTail(self.flags)
+
+    # useful ? here ?
+    @property     
+    def is_var(self):
+        return idc.isVar(self.flags)
+
+    @property
+    def has_extra_comment(self):
+        """ 
+            Does this address has extra prev ou next line comments ?
+             - see LineA and LineB
+        """
+        return idc.isExtra(self.flags)
+
+    @property
+    def has_ref(self):
+        return idc.isRef(self.flags) 
+
+    @property
+    def has_value(self):
+        return idc.hasValue(self.flags) 
+
+    # comments: property ? for normal and repeteable ?   
+    def set_comment(self, comment, repeteable=True):
+        if repeteable:
+            idc.MakeRptCmt(self.addr, comment)
+        else:
+            return idc.MakeComm(self.addr, comment)
+
+    def get_comment(self, repeteable=True):
+        return idc.CommentEx(self.addr, repeteable)
+
 
 class IDANamedElt(IDAElt):
     """ Real base class : looks like everything can have a name """
@@ -52,11 +119,25 @@ def set_name(self, name):
 
     def __IDA_repr__(self):
         if self.name is not "":
-            return self.__class__.__name__ + " " +  self.name
-        return super(IDANamedElt, self).__IDA_repr__()
+            return self.name
+        return   "{no name}"
+
+    # Do not use the Has*Name from idc because these have no sens
+    @property
+    def has_user_name(self):
+        return bool(self.flags & idc.FF_NAME)
+
+    @property    
+    def has_dummy_name(self):
+        return bool(self.flags & idc.FF_LABL)
+
+    @property  
+    def has_name(self):
+        return bool(self.flags & idc.FF_ANYNAME)
 
 
 class IDASizedElt(IDAElt):
+    # always use NextHead to get endaddr ?
     def __init__(self, addr, endaddr, nb_elt=None):
         """ endaddr: first addr not part of the element """
         super(IDASizedElt, self).__init__(addr, endaddr, nb_elt)
@@ -69,6 +150,30 @@ def __init__(self, addr, endaddr, nb_elt=None):
     def __contains__(self, value):
         return self.addr <= value < self.endADDR
 
+    def patch(self, patch, fill_nop=True):
+        print("PATCH ASKED at <{0}| size {1}> with {2}".format(self.addr, self.size, patch))
+        nop = 0x90 #<- need to adapt to other platform
+        if self.size < len(patch):
+            raise ValueError("Patch if too big for {0}".format(self)) 
+        if self.size != len(patch) and not fill_nop:
+            raise Value("Patch is too small for {0} and no fill_patch (better idea than raise ?)".format(self))
+
+        full_patch = list(patch) + [nop] * (self.size - len(patch))
+        for addr, byte in zip(range(self.addr, self.addr + self.size), full_patch):
+            if idc.Byte(addr) == byte:
+                print("NOPATCH BYTE : SAME VALUE")
+                continue
+            if not idc.PatchByte(addr, byte):
+                print("PATCH addr {0} with byte {1} failed".format(hex(addr), hex(byte)))
+
+    def replace(self, value):
+        return self.patch([value] * self.size)
+
+    @property
+    def bytes(self):
+        return [data.ByteData(addr) for addr in range(self.addr, self.addr + self.size)]
+
+
 
 class IDANamedSizedElt(IDASizedElt, IDANamedElt):
     pass