#! /usr/bin/env python3
""" A python script used to decode the addresses in traces generated by
the dotrace instrumentation.

Requires addr2line to be installed. The script works by collecting a
given number of address from the trace file and then calling addr2line
with such a block of addresses.
"""

import os
import subprocess
import collections
import sys
from itertools import chain
import argparse

def decode(bin, traces):
	"""Decodes a list of address using addr2line. 

	Args:
		bin: The executable used for retrieving the debug symbols.
		traces: A list of lists of the form [c, a, b] where a and 
		b are addresses as string.

	Returns:
		Generates lists [c, call targed name, call site name, 
		call site source location]
	"""
	if len(traces) == 0: 
		return []
	addrs = chain.from_iterable([x[1:] for x in traces])
	cmd = ["addr2line", "-f", "-s", "-e", bin] + list(addrs)
	decoded = subprocess.run(cmd, stdout=subprocess.PIPE).stdout 
	decoded = collections.deque(decoded.split(b'\n'))
	for line in traces:
		target_name = decoded.popleft()
		target_line = decoded.popleft()
		site_name = decoded.popleft()
		site_line = decoded.popleft()
		yield [line[0],target_name,site_name,site_line]

def processBlock(args, traces, out_file):
	"""Decodes a block of trace lines.

	Args:
		args: The program arguments.
		traces: A list of lists of the form [c, a, b] where a and 
		b are addresses as string and c is eiter 'e' or 'x'
		indicating if the line is a function entry or exit.
		out_file: The output file to write into. In binary mode.
	"""
	decoded = decode(args.binary, traces)
	for decoded_line in decoded:
		if decoded_line[0] == 'e':
			out_file.write(b"Enter ")
			out_file.write(decoded_line[1])
			if not args.suppress_callsite:
				out_file.write(b" from ")
				out_file.write(decoded_line[2])
				out_file.write(b" (")
				out_file.write(decoded_line[3])
				out_file.write(b")")
			out_file.write(b"\n")
		if decoded_line[0] == 'x' and (not args.suppress_exit):
			out_file.write(b"Exit ")
			out_file.write(decoded_line[1])
			out_file.write(b"\n")

# Parser configuration
parser = argparse.ArgumentParser(description='Decode a trace file generated by dotrace.')
parser.add_argument('-b', '--binary', help='The binary file.', 
	required=True, metavar="FILE")
parser.add_argument('-t', '--trace', help="The trace file (default: '%(default)s').", 
	default="trace.out", metavar="FILE")
parser.add_argument('-o', '--output', help="The output file (default: '%(default)s').", 
	default="trace.decoded", metavar="FILE")
parser.add_argument('-e', '--suppress-exit', 
	help='Do not output function exits.', action='store_true')
parser.add_argument('-c', '--suppress-callsite', 
	help='Do not output callsite of function calls.', action='store_true')
parser.add_argument('-s', '--block-size', type=int, default=10000,
	help='The number of addresses forwarded to addr2line (default: %(default)s).')

if __name__ == "__main__":
	args = parser.parse_args()

	# Check if addr2line is installed
	DEVNULL = open(os.devnull, 'w')
	try:
		subprocess.call(["addr2line", "--version"], 
			stdout=DEVNULL, stderr=subprocess.STDOUT)
	except FileNotFoundError as e:
		print("Error: addr2line not installed.", file=sys.stderr)
		sys.exit(1)

	# Check if the files referenced by the options exist
	for f in [args.binary, args.trace]:
		if not os.path.exists(f):
			print("Error: File {} does not exist.".format(f) ,file=sys.stderr)
			sys.exit(1)

	lines_to_read = args.block_size / 2
	accumulator = []
	i = 0
	total = 1
	with open(args.trace, 'r') as trace_file:
		with open(args.output, 'bw') as out_file:
			# read lines_to_read lines from the file, each line constains two
			# addresses
			for line in trace_file:
				i = i + 1
				accumulator.append(line.split())
				if i >= lines_to_read:
					print("Processing lines {}-{}.".format(total, total + i - 1))
					processBlock(args, accumulator,  out_file)
					accumulator.clear() 
					total = total + i
					i = 0
			print("Processing lines {}-{} [final block].".format(total, total + i - 1))
			processBlock(args, accumulator, out_file)