def _EncodeAny(t, f, dl, s, issueError): s = ut.unProxy(s) if isinstance(s, str) or isinstance(s, bytes): try: return f(s) except KeyError: if isinstance(s, str): try: return dl[s] except KeyError: sl = ["Cannot encode string %s as %s." % (s, t)] for match in difflib.get_close_matches(s, dl.keys()): sl.append(' - Suggestion: %s' % match) raise ut.EPError('\n'.join(sl)) if issueError: raise ut.EPError('[Warning] "%s" is not a %s' % (s, t)) return s else: try: return dl.get(s, s) except TypeError: # unhashable return s
def AddString(self, string): string = ut.u2b(string) # Starcraft uses multibyte encoding. if not isinstance(string, bytes): raise ut.EPError('Invalid type for string') stringindex = len(self._dataindextb) # If duplicate text exist -> just proxy it try: repr_stringid = self._stringmap[string] dataindex = self._dataindextb[repr_stringid] self._dataindextb.append(dataindex) self._capacity += 2 # just string offset # Else -> Create new entry except KeyError: dataindex = len(self._datatb) self._stringmap[string] = stringindex self._datatb.append(string) self._dataindextb.append(dataindex) # string + b'\0' + string offset self._capacity += len(string) + 1 + 2 ut.ep_assert(self._capacity < 65536, 'String table overflow') return stringindex
def CollectObjects(root): global phase global _rootobj global _found_objects global _found_objects_set global _dynamic_objects_set global _untraversed_objects lprint('[Stage 1/3] CollectObjects', flush=True) phase = PHASE_COLLECTING objc = ObjCollector() _rootobj = None _found_objects_set = set() _dynamic_objects_set = set() _untraversed_objects = [] # Evaluate root to register root object. # root may not have WritePayload() method e.g: Forward() constexpr.Evaluate(root) while _untraversed_objects: while _untraversed_objects: lprint(" - Collected %d / %d objects" % (len(_found_objects_set), len(_found_objects_set) + len(_untraversed_objects))) obj = _untraversed_objects.pop() objc.StartWrite() obj.CollectDependency(objc) objc.EndWrite() # Check for new objects for obj in list(_dynamic_objects_set): objc.StartWrite() obj.CollectDependency(objc) objc.EndWrite() if len(_found_objects_set) == 0: raise ut.EPError('No object collected') # Shuffle objects -> Randomize(?) addresses _found_objects_set.remove(_rootobj) fo2 = list(_found_objects_set) random.shuffle(fo2) _found_objects = [_rootobj] + fo2 # cleanup _found_objects_set = None phase = None # Final lprint(" - Collected %d / %d objects" % (len(_found_objects), len(_found_objects)), flush=True)
def GetStringIndex(self, string): string = ut.u2b(string) if not isinstance(string, bytes): raise ut.EPError('Invalid type for string') try: return self._stringmap[string] + 1 except KeyError: return self.AddString(string) + 1
def AllocObjects(): global phase global _alloctable global _payload_size phase = PHASE_ALLOCATING objn = len(_found_objects) lprint("[Stage 2/3] AllocObjects", flush=True) # Quick and less space-efficient approach if not _payload_compress: lallocaddr = 0 for i, obj in enumerate(_found_objects): objsize = obj.GetDataSize() allocsize = (objsize + 3) & ~3 _alloctable[obj] = lallocaddr lallocaddr += allocsize lprint(" - Allocated %d / %d objects" % (i + 1, objn)) _payload_size = lallocaddr lprint(" - Allocated %d / %d objects" % (objn, objn), flush=True) phase = None return obja = ObjAllocator() _alloctable = {} dwoccupmap_dict = {} # Get occupation map for all objects for i, obj in enumerate(_found_objects): obja.StartWrite() obj.WritePayload(obja) dwoccupmap = obja.EndWrite() dwoccupmap_dict[obj] = dwoccupmap if len(dwoccupmap) != (obj.GetDataSize() + 3) >> 2: raise ut.EPError( 'Occupation map length (%d) & Object size mismatch for object (%d)' % (len(dwoccupmap), (obj.GetDataSize() + 3) >> 2)) lprint(" - Preprocessed %d / %d objects" % (i + 1, objn)) lprint(" - Preprocessed %d / %d objects" % (objn, objn), flush=True) lprint(" - Allocating objects..", flush=True) stackobjs.StackObjects(_found_objects, dwoccupmap_dict, _alloctable) # Get payload length _payload_size = max( map(lambda obj: _alloctable[obj] + obj.GetDataSize(), _found_objects)) phase = None
def CompressPayload(mode): ''' Set payload compression mode. :param mode: If true, enable object stacking (compression). If false, disable it. ''' global _payload_compress if mode not in [True, False]: raise ut.EPError('Invalid type') if mode: _payload_compress = True else: _payload_compress = False
def WriteDword(self, number): number = constexpr.Evaluate(number) if number.rlocmode: ut.ep_assert( self._datacur % 4 == 0, 'Non-const dwords must be aligned to 4byte' ) if number.rlocmode == 1: self._prttable.append(self._datacur) elif number.rlocmode == 4: self._orttable.append(self._datacur) else: raise ut.EPError('rlocmode should be 1 or 4') offset = number.offset self._data[self._datacur + 0] = offset & 0xFF self._data[self._datacur + 1] = (offset >> 8) & 0xFF self._data[self._datacur + 2] = (offset >> 16) & 0xFF self._data[self._datacur + 3] = (offset >> 24) & 0xFF self._datacur += 4
def f_dbstr_print(dst, *args): """Print multiple string / number to dst. :param dst: Destination address (Not EPD player) :param args: Things to print """ if ut.isUnproxyInstance(dst, DBString): dst = dst.GetStringMemoryAddr() args = ut.FlattenList(args) for arg in args: if ut.isUnproxyInstance(arg, bytes): dst = f_dbstr_addstr(dst, c.Db(arg + b'\0')) elif ut.isUnproxyInstance(arg, str): dst = f_dbstr_addstr(dst, c.Db(ut.u2b(arg) + b'\0')) elif ut.isUnproxyInstance(arg, DBString): dst = f_dbstr_addstr(dst, arg.GetStringMemoryAddr()) elif ut.isUnproxyInstance(arg, int): # int and c.EUDVariable should act the same if possible. # EUDVariable has a value of 32bit unsigned integer. # So we adjust arg to be in the same range. dst = f_dbstr_addstr(dst, c.Db( ut.u2b(str(arg & 0xFFFFFFFF)) + b'\0')) elif ut.isUnproxyInstance(arg, c.EUDVariable): dst = f_dbstr_adddw(dst, arg) elif c.IsConstExpr(arg): dst = f_dbstr_adddw(dst, arg) elif ut.isUnproxyInstance(arg, hptr): dst = f_dbstr_addptr(dst, arg._value) else: raise ut.EPError( 'Object wit unknown parameter type %s given to f_eudprint.' % type(arg) ) return dst
def WritePayload(self, pbuffer): """Write object""" raise ut.EPError('Override')
def GetDataSize(self): """Memory size of object.""" raise ut.EPError('Override')
def _GetLastLoopBlock(): for block in reversed(ut.EUDGetBlockList()): if _IsLoopBlock(block[1]): return block raise ut.EPError('No loop block surrounding this code area')
def getValueAddr(self): raise ut.EPError('override')