diff --git a/ebpfcat/arraymap.py b/ebpfcat/arraymap.py index bc63db6d6aa3db8003917b28dade1380015bc7de..2387c86f00b0101992f1a62cff199c0daa669b5a 100644 --- a/ebpfcat/arraymap.py +++ b/ebpfcat/arraymap.py @@ -1,39 +1,39 @@ -from struct import pack, unpack, calcsize +from struct import pack_into, unpack_from, calcsize -from .ebpf import FuncId, Map, Memory, Opcode, SubProgram +from .ebpf import FuncId, Map, Memory, MemoryDesc, Opcode from .bpf import create_map, lookup_elem, MapType, update_elem -class ArrayGlobalVarDesc: - def __init__(self, map, size): +class ArrayGlobalVarDesc(MemoryDesc): + base_register = 0 + + def __init__(self, map, fmt): self.map = map - self.fmt = size + self.fmt = fmt - def __get__(self, ebpf, owner): - if ebpf is None: - return self - position = ebpf.__dict__[self.name] - if isinstance(ebpf, SubProgram): - ebpf = ebpf.ebpf - if ebpf.loaded: - data = ebpf.__dict__[self.map.name].data[ - position : position+calcsize(self.fmt)] - return unpack(self.fmt, data)[0] - return Memory(ebpf, Memory.fmt_to_opcode[self.fmt], - ebpf.r0 + position, self.fmt.islower()) + def fmt_addr(self, ebpf): + return self.fmt, ebpf.__dict__[self.name] def __set_name__(self, owner, name): self.name = name - def __set__(self, ebpf, value): - position = ebpf.__dict__[self.name] - if isinstance(ebpf, SubProgram): - ebpf = ebpf.ebpf - if ebpf.loaded: - ebpf.__dict__[self.map.name].data[ - position : position+calcsize(self.fmt)] = pack(self.fmt, value) + def __get__(self, instance, owner): + if instance is None: + return self + fmt, addr = self.fmt_addr(instance) + if instance.ebpf.loaded: + data = instance.ebpf.__dict__[self.map.name].data + return unpack_from(fmt, data, addr)[0] + else: + return super().__get__(instance, owner) + + def __set__(self, instance, value): + fmt, addr = self.fmt_addr(instance) + if instance.ebpf.loaded: + pack_into(fmt, instance.ebpf.__dict__[self.map.name].data, + addr, value) else: - getattr(ebpf, f"m{self.fmt}")[ebpf.r0 + position] = value + super().__set__(instance, value) class ArrayMapAccess: @@ -49,8 +49,8 @@ class ArrayMapAccess: class ArrayMap(Map): - def globalVar(self, signed=False, size=4): - return ArrayGlobalVarDesc(self, size, signed) + def globalVar(self, fmt="I"): + return ArrayGlobalVarDesc(self, fmt) def add_program(self, owner, prog): position = getattr(owner, self.name) diff --git a/ebpfcat/ebpf.py b/ebpfcat/ebpf.py index d523cc989196d46c9b428ddcf63d169c0ab63918..a047365c3a07d881e35d4adbe926b0460d5f028e 100644 --- a/ebpfcat/ebpf.py +++ b/ebpfcat/ebpf.py @@ -664,24 +664,18 @@ class MemoryDesc: def __get__(self, instance, owner): if instance is None: return self - elif isinstance(instance, SubProgram): - ebpf = instance.ebpf - else: - ebpf = instance fmt, addr = self.fmt_addr(instance) - return Memory(ebpf, Memory.fmt_to_opcode[fmt], - ebpf.r[self.base_register] + addr, fmt.islower()) + return Memory(instance.ebpf, Memory.fmt_to_opcode[fmt], + instance.ebpf.r[self.base_register] + addr, + fmt.islower()) def __set__(self, instance, value): - if isinstance(instance, SubProgram): - ebpf = instance.ebpf - else: - ebpf = instance + ebpf = instance.ebpf fmt, addr = self.fmt_addr(instance) bits = Memory.fmt_to_opcode[fmt] if isinstance(value, int): ebpf.append(Opcode.ST + bits, self.base_register, 0, - self.addr(instance), value) + addr, value) return elif isinstance(value, IAdd): value = value.value @@ -987,6 +981,10 @@ class EBPF: yield self.stack self.stack = oldstack + @property + def ebpf(self): + return self + tmp = TemporaryDesc(None, "r") stmp = TemporaryDesc(None, "sr") wtmp = TemporaryDesc(None, "w")