#!/usr/bin/env python
#
# Copyright 2014 John-Mark Gurney.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#

import struct
import weakref

from ctypes import CDLL, Structure, POINTER, pointer, create_string_buffer, sizeof, CFUNCTYPE
from ctypes import c_int, c_uint, c_char_p, c_long, c_ulong, c_size_t, c_ssize_t, c_void_p

_ctflib = CDLL('libctf.so.2')
#_ctflib = CDLL('/usr/obj/arm.armeb/usr/src.avila/tmp/usr/src.avila/cddl/lib/libctf/libctf.so.2')

class CTFError(Exception):
	pass

class HexInt(long):
	def __str__(self):
		return '%#x' % self

	__repr__ = __str__

models = {
	'CTF_MODEL_ILP32':	1,
	'CTF_MODEL_LP64':	2,
}
models.update((v, k) for k, v in models.items())

# From: sys/ctf.h
kinds = {
	'CTF_K_UNKNOWN':	0,
	'CTF_K_INTEGER':	1,
	'CTF_K_FLOAT':		2,
	'CTF_K_POINTER':	3,
	'CTF_K_ARRAY':		4,
	'CTF_K_FUNCTION':	5,
	'CTF_K_STRUCT':		6,
	'CTF_K_UNION':		7,
	'CTF_K_ENUM':		8,
	'CTF_K_FORWARD':	9,
	'CTF_K_TYPEDEF':	10,
	'CTF_K_VOLATILE':	11,
	'CTF_K_CONST':		12,
	'CTF_K_RESTRICT':	13,
}
kinds.update((v, k) for k, v in kinds.items())

intenc = {
	'CTF_INT_SIGNED':	0x01,
	'CTF_INT_CHAR':		0x02,
	'CTF_INT_BOOL':		0x04,
	'CTF_INT_VARARGS':	0x08,
}
intenc.update((v, k) for k, v in intenc.items())

def lookupid(v, d):
	tv = d[v]
	try:
		return int(v)
	except ValueError:
		return tv

def lookupstr(v, d):
	tv = d[v]
	if isinstance(tv, str):
		return tv
	else:
		return d[tv]

class ctf_visit_f(Structure):
	pass

class ctf_encoding_t(Structure):
	_fields_ = [ ('cte_format', c_uint),
		('cte_offset', c_uint),
		('cte_bits', c_uint),
		]

ctf_id_t = c_long

# If c_int != 0, return immediately
ctf_member_f = CFUNCTYPE(c_int, c_char_p, ctf_id_t, c_ulong, c_void_p)

class ctf_file_t(Structure):
	pass

ctf_file_t_p = POINTER(ctf_file_t)

for r, f, a in (
	(ctf_file_t_p,	'ctf_open',		(c_char_p, POINTER(c_int), )),
	(None,		'ctf_close',		(ctf_file_t_p, )),
	(c_int,		'ctf_update',		(ctf_file_t_p, )),
	(c_int,		'ctf_discard',		(ctf_file_t_p, )),
	(c_int,		'ctf_write',		(ctf_file_t_p, c_int, )),

	(ctf_file_t_p,	'ctf_parent_file',	(ctf_file_t_p, )),
	(c_char_p,	'ctf_parent_name',	(ctf_file_t_p, )),

	(c_int,		'ctf_setmodel',		(ctf_file_t_p, c_int, )),
	(c_int,		'ctf_getmodel',		(ctf_file_t_p, )),

	(c_int,		'ctf_errno',		(ctf_file_t_p, )),
	(c_char_p,	'ctf_errmsg',		(c_int, )),
	(c_int,		'ctf_version',		(c_int, )),

	(ctf_id_t,	'ctf_lookup_by_name',	(ctf_file_t_p, c_char_p, )),
	(ctf_id_t,	'ctf_type_resolve',	(ctf_file_t_p, ctf_id_t, )),
	(c_ssize_t,	'ctf_type_lname',	(ctf_file_t_p, ctf_id_t, c_char_p, c_size_t, )),
	(c_char_p,	'ctf_type_name',	(ctf_file_t_p, ctf_id_t, c_char_p, c_size_t, )),
	(c_ssize_t,	'ctf_type_size',	(ctf_file_t_p, ctf_id_t, )),
	(c_ssize_t,	'ctf_type_align',	(ctf_file_t_p, ctf_id_t, )),
	(int,		'ctf_type_kind',	(ctf_file_t_p, ctf_id_t, )),
	(ctf_id_t,	'ctf_type_reference',	(ctf_file_t_p, ctf_id_t, )),
	(ctf_id_t,	'ctf_type_pointer',	(ctf_file_t_p, ctf_id_t, )),
	(int,		'ctf_type_encoding',	(ctf_file_t_p, ctf_id_t, POINTER(ctf_encoding_t), )),
	(int,		'ctf_type_visit',	(ctf_file_t_p, ctf_id_t, POINTER(ctf_visit_f), c_void_p, )),
	(int,		'ctf_type_cmp',		(ctf_file_t_p, ctf_id_t, ctf_file_t_p, ctf_id_t, )),
	(int,		'ctf_type_compat',	(ctf_file_t_p, ctf_id_t, ctf_file_t_p, ctf_id_t, )),
    ):
	fun = getattr(_ctflib, f)
	setattr(fun, 'restype', r)
	setattr(fun, 'argtypes', a)

class CTFID:
	def __init__(self, _ctfobj, id):
		self._ctfobj = _ctfobj
		self._ctf = _ctfobj._ctf
		self._id = id

	__int_sizes = {
		1: 'B',
		2: 'H',
		4: 'I',
		8: 'Q',
	}

	def __call__(self, s, endian='@'):
		size = self.getsize()

		if len(s) < size:
			raise ValueError('not enough bytes, needed %d', size)

		base = self.resolve()
		kind = base.getkind()

		src, rem = s[:size], s[size:]

		conv = endian
		if kind == 'CTF_K_INTEGER':
			enc = base.getencoding()
			if enc.cte_format & intenc['CTF_INT_CHAR']:
				assert size == 1
				val = src
			else:
				convfun = lambda x: x
				if enc.cte_format & intenc['CTF_INT_SIGNED']:
					convfun = lambda x: x.lower()

				if enc.cte_format & intenc['CTF_INT_BOOL']:
					convchar = '?'
				else:
					convchar = self.__int_sizes[size]

				conv += convfun(convchar)
				val = struct.unpack(conv, src)[0]
		elif kind == 'CTF_K_POINTER':
			val = HexInt(struct.unpack(endian + self.__int_sizes[size], src)[0])
		elif kind == 'CTF_K_ENUM':
			raise NotImplementedError(kind)
		elif kind in ('CTF_K_STRUCT', 'CTF_K_UNION'):
			members = []
			def addmember(*args):
				members.append(args)
				return 0

			self.memberiter(addmember)

			#print `kind`
			#print `members`
			val = {}
			for name, _type, off in members:
				if off % 8 != 0:
					raise NotImplementedError('bytes only for now')
				off /= 8
				#print 'h:', `src`, `off`, `_type.getsize()`
				val[name], trem = _type(src[off:off + _type.getsize()], endian)
				assert not trem, 'trem left? %s' % `trem`
		elif kind == 'CTF_K_UNION':
			raise NotImplementedError(kind)
		else:
			raise NotImplementedError(kind)

		return val, rem

	def memberiter(self, fun):
		iterfun = self._ctfobj.makeiterfun(fun)
		_ctflib.ctf_member_iter(self._ctf, self._id, iterfun, None)

	def resolve(self):
		return self._ctfobj.resolve(self._id)

	def reference(self):
		return self._ctfobj.reference(self._id)

	def getname(self):
		buf = create_string_buffer(64)
		rbuf = _ctflib.ctf_type_name(self._ctf, self._id, buf, sizeof(buf))
		return rbuf

	def getalign(self):
		return _ctflib.ctf_type_align(self._ctf, self._id)

	def getsize(self):
		return _ctflib.ctf_type_size(self._ctf, self._id)

	def getkind(self):
		return lookupstr(_ctflib.ctf_type_kind(self._ctf, self._id),
		    kinds)

	def getencoding(self):
		enc = ctf_encoding_t()
		r = _ctflib.ctf_type_encoding(self._ctf, self._id, pointer(enc))
		if r != 0:
			raise self._ctfobj.geterror()

		return enc

	def __repr__(self):
		i = map(lambda x: ': '.join(x), [
			('name', str(self.getname())),
			('size', str(self.getsize())),
			('kind', self.getkind()),
			('align', str(self.getalign())),
		])
		return '<CTFID: %s>' % ', '.join(i)

class CTF:
	def __init__(self, f, model=None):
		i = c_int()
		self._ctf = _ctflib.ctf_open(f, pointer(i))
		if not bool(self._ctf):
			raise ValueError('%s: %s' % (`f`,
			    _ctflib.ctf_errmsg(i)))

		self._ids = weakref.WeakValueDictionary()

		if model is not None:
			_ctflib.ctf_setmodel(self._ctf, lookupid(model, models))

	def makeiterfun(self, fun):
		def lfun(name, id, off, ptr):
			id = self._getid(id)
			return fun(name, id, off)

		return ctf_member_f(lfun)

	def resolve(self, id):
		if isinstance(id, CTFID):
			id = id._id

		nid = _ctflib.ctf_type_resolve(self._ctf, id)
		return self._getid(nid)

	def reference(self, id):
		if isinstance(id, CTFID):
			id = id._id

		nid = _ctflib.ctf_type_reference(self._ctf, id)
		return self._getid(nid)

	def geterror(self):
		err = _ctflib.ctf_errno(self._ctf)
		errmsg = _ctflib.ctf_errmsg(err)
		return CTFError((err, errmsg))

	def __getitem__(self, k):
		id = _ctflib.ctf_lookup_by_name(self._ctf, k)
		if id == -1:
			raise KeyError('unknown name: %s' % `k`)
		#print 'getitem:', `id`
		return self._getid(id)

	def _getid(self, id):
		try:
			return self._ids[id]
		except KeyError:
			obj = CTFID(self, id)
			self._ids[id] = obj

			return obj

	def __del__(self):
		try:
			_ctflib.ctf_close(self._ctf)
		except Exception:
			pass

def dump(name, _type, off):
	_type = _type.resolve()
	print 'd:', `name, _type, off`
	return 0

def parseddb(s):
	listofchars = [ x.split(':', 1)[1].split() for x in s.split('\n') ]

	chars = sum(listofchars, [])
	#print chars
	l = ''.join(chr(int(x, 16)) for x in chars)

	return l

if __name__ == '__main__':
	#ctf = CTF('/boot/kernel/kernel')
	ctf = CTF('/tmp/avila3.ctf', model='CTF_MODEL_ILP32')
	#ctf = CTF('/tftpboot/kernel.avila.avila')

	i = ctf['int']
	vm = ctf['vm_page_t']
	vmr = vm.resolve()
	vmrr = vmr.reference()
	print 'int:', `i`
	print 'vm_page_t:', `vm`
	print 'vm_page_t res:', `vmr`
	print 'vm_page_t resolve:', `vmrr`
	#vmrr.memberiter(dump)

	#print `vmrr('\x11' * vmrr.getsize())`
	r = parseddb('''0xc0805db0:     0   0   0   0   c0  6d  1f  64  c0  80  5e  0   c0  6d  18  a0
0xc0805dc0:     c0  6d  18  78  0   0   0   0   0   0   0   0   0   0   d   f4
0xc0805dd0:     1   f   40  0   0   0   0   0   0   0   0   0   c0  df  40  0
0xc0805de0:     0   0   0   0   c0  80  5d  e0  0   0   0   1   0   0   0   1
0xc0805df0:     0   0   0   0   0   4   ff  1   9   0   0   ff  ff  0   0   0
0xc0805e00:     0   0   0   0   c0  6d  1f  4c  c0  80  5e  50  c0  80  5d  b8
0xc0805e10:     c0  6d  18  78  0   0   0   0   0   0   0   0   0   0   d   f5
0xc0805e20:     1   f   50  0   0   0   0   0   0   0   0   0   c0  df  50  0''')
	print len(r), `r`
	import pprint
	s, r = vmrr(r, '>')
	pprint.pprint(s)

	vm = ctf['vm_object_t'].resolve()
	print `vm`