1#!/usr/bin/env python3
2##
3## This file is part of the sigrok-util project.
4##
5## Copyright (C) 2012 Bert Vermeulen <bert@biot.com>
6##
7## This program is free software; you can redistribute it and/or modify
8## it under the terms of the GNU General Public License as published by
9## the Free Software Foundation; either version 3 of the License, or
10## (at your option) any later version.
11##
12## This program is distributed in the hope that it will be useful,
13## but WITHOUT ANY WARRANTY; without even the implied warranty of
14## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15## GNU General Public License for more details.
16##
17## You should have received a copy of the GNU General Public License
18## along with this program; if not, see <http://www.gnu.org/licenses/>.
19##
20
21import sys
22import os
23from getopt import getopt
24import struct
25
26
27def parse(filename):
28	f = open(filename, 'rb')
29	if f.read(2) != b'MZ':
30		raise Exception("MZ signature not found.")
31
32	sections = []
33	# long e_lfanew
34	f.seek(0x3c)
35	pe_ptr = struct.unpack("<L", f.read(4))[0]
36	f.seek(pe_ptr)
37	if f.read(4) != b'\x50\x45\x00\x00':
38		raise Exception("PE signature not found.")
39	# skip Machine
40	f.seek(f.tell() + 2)
41	sections.append(['header', 324, 0])
42	num_sections = struct.unpack("<H", f.read(2))[0]
43	# skip TimeDateStamp
44	f.seek(f.tell() + 4)
45	symboltable_address = struct.unpack("<L", f.read(4))[0]
46	num_symbols = struct.unpack("<L", f.read(4))[0]
47	optheader_size = struct.unpack("<H", f.read(2))[0]
48	# skip past PE header and PE optional header
49	f.seek(f.tell() + 2 + optheader_size)
50
51	for i in range(num_sections):
52		name = f.read(8).decode('ascii', errors='ignore').strip('\x00')
53		# skip past Misc and VirtualAddress
54		f.seek(f.tell() + 8)
55		# SizeOfRawData
56		size = struct.unpack("<L", f.read(4))[0]
57		# PointerToRawData
58		ptr = struct.unpack("<L", f.read(4))[0]
59		# skip to next section header
60		f.seek(f.tell() + 16)
61		sections.append([name, size, ptr])
62
63	symbols = []
64	addr = symboltable_address
65	stringtable_address = symboltable_address + num_symbols * 18
66	for i in range(num_symbols):
67		f.seek(addr)
68		tmp = f.read(8)
69		symaddr = struct.unpack("<L", f.read(4))[0]
70		# skip SectionNumber and Type
71		symtype = struct.unpack("B", f.read(1))[0]
72		f.seek(f.tell() + 4)
73		if tmp[:4] == b'\x00\x00\x00\x00':
74			# symbol name is in the string table
75			straddr = stringtable_address + struct.unpack("<l", tmp[4:])[0]
76			f.seek(straddr)
77			tmpname = f.read(64)
78			name = tmpname[:tmpname.find(b'\x00')]
79		else:
80			name = tmp
81		name = name.decode('ascii', errors='ignore').strip('\x00')
82		# need IMAGE_SYM_CLASS_EXTERNAL
83		if symtype == 0x02:
84			size = 0
85		else:
86			size = None
87		if i != 0 and symbols[-1][2] is not None and symaddr > symbols[-1][1]:
88			symbols[-1][2] = symaddr - symbols[-1][1]
89		symbols.append([name, symaddr, size])
90		addr += 18
91
92	f.close()
93
94	return sections, symbols
95
96
97def list_all(filename):
98	sections, symbols = parse(filename)
99	if sections:
100		print("Sections:\n    Name         Size\t  Position")
101		cnt = 0
102		for name, size, address in sections:
103			print("%-3d %-8s    %5d\t  0x%.8x" % (cnt, name, size, address))
104			cnt += 1
105	if symbols:
106		print("\nSymbol table:\n   Address     Size    Symbol")
107		for symbol, address, size in symbols:
108			if size is not None:
109				sizestr = "%5d" % size
110			else:
111				sizestr = "     "
112			print("0x%.8x  %s    %-8s" % (address, sizestr, symbol))
113		print()
114
115
116def extract_symbol(filename, symbol):
117	sections, symbols = parse(filename)
118	if not symbols:
119		return None
120	data = None
121	for symbolname, address, size in symbols:
122		if symbolname == symbol:
123			if size is None:
124				raise Exception("symbol %s found, but has unknown size")
125			f = open(filename, 'rb')
126			f.seek(address)
127			data = f.read(size)
128			f.close()
129			if len(data) != size:
130				raise Exception("short file")
131			break
132
133	if data is None:
134		raise Exception("symbol %s not found" % symbol)
135
136	return data
137
138
139
140def usage():
141	print("usage: parsepe.py [-s <symbol>] <-l|-x> <filename>")
142	print("  -l   list all sections and symbols in file")
143	print("  -x   extract symbol from file (specify symbol name with -s)")
144	sys.exit()
145
146
147#
148# main
149#
150
151if __name__ == '__main__':
152	filename = symbol = mode = None
153	opts, args = getopt(sys.argv[1:], 's:lx')
154	for opt, arg in opts:
155		if opt == '-s':
156			symbol = arg
157		elif opt == '-l':
158			mode = 'list'
159		elif opt == '-x':
160			mode = 'extract'
161
162	if len(args) != 1:
163		usage()
164	if mode is None and symbol is None:
165		usage()
166
167	try:
168		filename = args[0]
169		if mode == 'list':
170			list_all(filename)
171		elif mode == 'extract':
172			if symbol is None:
173				raise Exception("specify a symbol to extract")
174			data = extract_symbol(filename, symbol)
175			outfile = os.path.splitext(filename)[0] + symbol
176			open(outfile, 'wb').write(data)
177			print("saved %d bytes to %s" % (len(data), outfile))
178		else:
179			raise Exception("specify -l or -x")
180	except Exception as e:
181		print("Error: %s" % str(e))
182
183
184