bap.rpc

1 #!/usr/bin/env python 2 # -*- coding: utf-8 -*- 3 4 import os, time, atexit 5 from signal import signal, SIGTERM 6 import requests 7 from subprocess import Popen 8 from mmap import mmap 9 from urlparse import urlparse, parse_qs 10 from tempfile import NamedTemporaryFile 11 import json 12 import adt, arm, asm, bil 13 14 import threading 15 16 from pprint import pprint 17 18 __all__ = ["disasm", "image"] 19 20 DEBUG_LEVEL = ["Critical", "Error"] 21 22 storage = threading.local() 23 servers = dict() 24 server_lock = threading.Lock() 25 requests_lock = threading.Lock() 26 request = None 27

28 -def init_requests():

29 global request 30 with requests_lock: 31 if request == None: 32 request = requests.Session() 33 adapter = requests.adapters.HTTPAdapter( 34 pool_connections=1000, 35 pool_maxsize=1000, 36 max_retries=10, 37 pool_block=True) 38 request.mount('http://', adapter)

39 40 init_requests() 41 42

43 -def del_instance():

44 instance = getattr(storage, 'instance', None) 45 if instance is not None: 46 instance.close()

47

48 -def get_instance(**kwargs):

49 instance = getattr(storage, 'instance', None) 50 if instance is None: 51 args = kwargs.get('server', {}) 52 storage.instance = Bap(args) 53 return storage.instance

54 55 atexit.register(del_instance) 56 signal(SIGTERM, lambda x,y: del_instance) 57 58

59 -def spawn_server(**kwargs):

60 port = str(kwargs.get('port', 8080)) 61 name = kwargs.get('name', 'bap-server') 62 with server_lock: 63 if port in servers: 64 return servers[port] 65 else: 66 process = Popen([name, '--port=' + port]) 67 server = { 68 'server' : process, 69 'url' : "http://127.0.0.1:{0}".format(port) 70 } 71 servers[port] = server 72 return server

73 74

75 -def disasm(obj, **kwargs):

76 r""" disasm(obj) disassembles provided object. 77 Returns a generator object yielding instructions. 78 """ 79 def run(obj): 80 return get_instance(**kwargs).insns(obj, **kwargs)

81 if isinstance(obj, Id): 82 return run(obj) 83 elif isinstance(obj, Resource): 84 return run(obj.ident) 85 else: 86 return run(load_chunk(obj, **kwargs)) 87

88 -def image(f, **kwargs):

89 bap = get_instance(**kwargs) 90 if os.path.isfile(f) and not os.path.isabs(f): 91 f = os.path.abspath(f) 92 return Image(bap.load_file(f), bap)

93

94 -def load_chunk(s, **kwargs):

95 return get_instance(**kwargs).load_chunk(s, **kwargs)

96 97 98 99

100 -class Resource(object):

101 - def __init__(self, name, ident, bap):

102 self.ident = Id(ident) 103 self.bap = bap 104 self.msg = None 105 self._name = name

106

107 - def load(self):

108 if self.msg is None: 109 self.msg = self.bap.get_resource(self.ident) 110 if not self._name in self.msg: 111 if 'error' in msg: 112 raise ServerError(response) 113 else: 114 msg = "Expected {0} msg but got {1}".format( 115 self._name, msg) 116 raise RuntimeError(msg)

117

118 - def get(self, child):

119 self.load() 120 return self.msg[self._name].get(child)

121 122

123 -class Project(Resource):

124 - def __init__(self, ident, bap):

125 super(Image,self).__init__('program', ident, bap)

126

127 - def load_program(self):

128 self.program = bir.loads(self.get('program'))

129

130 - def __getattr__(self,name):

131 if name == 'program': 132 self.load_program() 133 return self.program 134 else: 135 return self.get(name)

136 137

138 -class Image(Resource):

139 - def __init__(self, ident, bap):

140 super(Image,self).__init__('image', ident, bap)

141

142 - def load_segments(self):

143 ss = self.get('segments') 144 self.segments = [Segment(s, self) for s in ss]

145

146 - def get_symbol(self, name, d=None):

147 for sec in self.segments: 148 sym = sec.get_symbol(name, d) 149 if sym is not d: 150 return sym 151 return d

152

153 - def __getattr__(self, name):

154 if name == 'segments': 155 self.load_segments() 156 return self.segments 157 else: 158 return self.get(name)

159

160 -class Segment(Resource):

161 - def __init__(self, ident, parent):

162 super(Segment, self).__init__('segment', ident, parent.bap) 163 self.parent = parent

164

165 - def load_symbols(self):

166 self.symbols = [Symbol(s, self) for s in self.get('symbols')]

167

168 - def get_symbol(self, name, d=None):

169 try: 170 return (s for s in self.symbols if s.name == name).next() 171 except StopIteration: 172 return d

173

174 - def __getattr__(self, name):

175 if name == 'symbols': 176 self.load_symbols() 177 return self.symbols 178 elif name == 'addr' or name == 'size': 179 return self.get('memory')[name] 180 elif name == 'memory': 181 self.memory = Memory(self.get('memory'), self) 182 return self.memory 183 else: 184 return self.get(name)

185

186 -class Symbol(Resource):

187 - def __init__(self, ident, parent):

188 super(Symbol, self).__init__('symbol', ident, parent.bap) 189 self.parent = parent

190

191 - def load_chunks(self):

192 self.chunks = [Memory(s, self) for s in self.get('chunks')]

193

194 - def __getattr__(self, name):

195 if name == 'chunks': 196 self.load_chunks() 197 return self.chunks 198 elif name == 'addr': 199 self.load_chunks() 200 return self.chunks[0].addr 201 else: 202 return self.get(name)

203

204 -class Memory(object):

205 - def __init__(self, mem, parent):

206 self.parent = parent 207 self.size = int(mem['size']) 208 self.addr = int(mem['addr']) 209 self.links = mem['links']

210

211 - def load_data(self):

212 try: 213 url = (urlparse(url) for url in self.links 214 if urlparse(url).scheme == 'mmap').next() 215 qs = parse_qs(url.query) 216 offset = int(qs['offset'][0]) 217 with open(url.path, "rw+b") as f: 218 mm = mmap(f.fileno(), length=0) 219 mm.seek(offset) 220 self.data = mm.read(self.size) 221 mm.close() 222 except StopIteration: 223 self.data = None

224

225 - def __getattr__(self, name):

226 if name == 'data': 227 self.load_data() 228 return self.data 229 raise AttributeError(name)

230 231

232 -class ServerError(Exception):

233 - def __init__(self, err):

234 self.msg = str(Error(err))

235

236 - def __str__(self):

237 return self.msg

238

239 -class Error(object):

240 - def __init__(self, err):

241 self.__dict__.update(err) 242 self.__dict__.update(err['error'])

243

244 - def __str__(self):

245 return "{severity}: {description}".format(**self.error)

246

247 -class Id(object):

248 - def __init__(self, r):

249 self.value = r

250 - def __str__(self):

251 return str(self.value)

252 253 RETRIES = 10 254

255 -class Bap(object):

256 - def __init__(self, server={}):

257 if isinstance(server, dict): 258 self.__dict__.update(spawn_server(**server)) 259 else: 260 self.url = server 261 262 self.last_id = 0 263 for attempt in range(RETRIES): 264 try: 265 self.capabilities = self.call({'init' : { 266 'version' : '0.1'}}).next()['capabilities'] 267 break 268 except Exception: 269 if attempt + 1 == RETRIES: 270 raise 271 else: 272 time.sleep(0.1 * attempt) 273 274 if not "capabilities" in self.__dict__: 275 raise RuntimeError("Failed to connect to BAP server") 276 self.data = {} 277 self.temp = NamedTemporaryFile('rw+b', prefix="bap-")

278

279 - def insns(self, src, **kwargs):

280 req = {'resource' : src} 281 req.update(kwargs) 282 res = self.call({'get_insns' : req}) 283 for msg in res: 284 if 'error' in msg: 285 err = Error(msg) 286 if err.severity in DEBUG_LEVEL: 287 print(err) 288 else: 289 return (parse_insn(js) for js in msg['insns'])

290

291 - def close(self):

292 self.__exit__()

293

294 - def load_file(self, name):

295 return self._load_resource({'load_file' : { 296 'url' : 'file://' + name}})

297

298 - def get_resource(self, name):

299 return self.call({'get_resource' : name}).next()

300

301 - def load_chunk(self, data, **kwargs):

302 kwargs.setdefault('url', self.mmap(data)) 303 kwargs.setdefault('arch', 'i386') 304 kwargs.setdefault('addr', 0) 305 addr = kwargs['addr'] 306 if isinstance(addr, str): 307 addr = int(addr, 0) 308 kwargs['addr'] = '0x{0:x}'.format(addr) 309 310 return self._load_resource({'load_memory_chunk' : kwargs})

311

312 - def __exit__(self):

313 if 'server' in self.__dict__: 314 self.server.terminate() 315 self.temp.close()

316

317 - def dumps(self,dic):

318 self.last_id += 1 319 dic['id'] = Id(self.last_id) 320 return json.dumps(dic, default=str)

321

322 - def call(self, data):

323 if isinstance(data, dict): 324 return jsons(request.post(self.url, data=self.dumps(data))) 325 else: 326 gen = (self.dumps(msg) for msg in data) 327 return jsons(request.post(self.url, data=gen))

328 329

330 - def mmap(self, data):

331 url = "mmap://{0}?offset=0&length={1}".format( 332 self.temp.name, len(data)) 333 os.ftruncate(self.temp.fileno(), 4096) 334 mm = mmap(self.temp.fileno(), 4096) 335 mm.write(data) 336 mm.close() 337 return url

338

339 - def _load_resource(self, res):

340 rep = self.call(res).next() 341 if 'error' in rep: 342 raise ServerError(rep) 343 return Id(rep['resource'])

344 345

346 -def jsons(r, p=0):

347 dec = json.JSONDecoder(encoding='utf-8') 348 while True: 349 obj,p = dec.scan_once(r.text,p) 350 yield obj

351

352 -def parse_target(js):

353 if 'target' in js: 354 return arm.loads(js['target']) 355 else: 356 return None

357

358 -def parse_bil(js):

359 if 'bil' in js: 360 return [bil.loads(s) for s in js['bil']] 361 else: 362 return None

363

364 -def parse_insn(js):

365 js.update(js['memory'], bil=parse_bil(js), target=parse_target(js)) 366 return asm.Insn(**js)

367

368 -def hexs(data):

369 return ' '.join(x.encode('hex') for x in data)

370

Source Code for Module bap.rpc