#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
filename:
	scanmbr.py
author:
	0xGiddi
date:
	18 Nov 2017 Sat 18 Nov 02:42:02 IST 2017
description:
	Search a specified MS Windows file and find (and extract)
	the code that belongs to boot sector (MBR). This is used
	for "fixing" or "rebuilding" a disk to boot MS Windows.
example:
	./scanmbr.py "/mnt/windows 7 iso/sources/winsetup.dll" "/tmp/msmbr.bin"

TODO:
	- !!! Rewrite the whole thing in a proper way... !!!!
	- Detect multiple signatures in single file
	- Read file in chunks (for large files)
	- Fix whole saving thingy
	- Add Windows 8/10 MBR detection (the MBR is slightly different)
"""

import sys
import re

MBR_SIGNATURES = [
	{'signature' : re.compile('\x33\xC0\x8E\xD0.{442}\x00{64}\x55\xAA'),
	 'description' : 'Windows 7 MBR (Empty partition table)',
	 'codesize'  : 446},
]


def main(inf, outf):
	try:
		# This whole thing needs to be rewritten...
	 	with open(inf, 'rb') as infile:
		 	data = infile.read()
		 	for signature in MBR_SIGNATURES:
		 		res = signature['signature'].search(data)
		 		if res:
		 			print "Signature '{desc}' found @ '0x{offset:02X}'".format(desc=signature['description'], offset=res.start())
		 			if outf:
		 				with open(outf, 'wb') as outfile:
		 					outfile.write(res.group()[:signature['codesize']])
		 			break
	except IOError as ex:
		print "Error parsing file ({err}, {num}".format(err=ex.message, num=ex.errno)



if __name__ == '__main__':
	if not sys.argv[1]:
		print "Usage: findmbr.py <inputfile> <outputfile>"
	outfile = None if len(sys.argv) == 2 else sys.argv[2]
	main(sys.argv[1], outfile)