mirror of
https://github.com/86Box/bios-tools.git
synced 2026-02-22 01:25:34 -07:00
275 lines
7.3 KiB
Python
275 lines
7.3 KiB
Python
import getopt, os, queue, subprocess, sys, threading, time
|
|
|
|
algos = ['lh5', 'lzari', 'bcd6f1', 'lzh']
|
|
longest_algo = 0
|
|
thread_status = []
|
|
term_size = os.get_terminal_size()
|
|
result_line = 1
|
|
max_result_line = 1
|
|
longest_result_line = 0
|
|
info_offset = 0
|
|
info_offset_found = 0
|
|
|
|
print_lock = threading.Lock()
|
|
def print_(s):
|
|
with print_lock:
|
|
sys.stdout.write(s)
|
|
sys.stdout.flush()
|
|
|
|
def print_thread():
|
|
global info_offset
|
|
global thread_status
|
|
|
|
# Initialize thread status display.
|
|
for thread_id in range(len(thread_status)):
|
|
sys.stdout.write('\033[{0};1H{0}'.format(thread_id + 1))
|
|
|
|
# Print thread status in a loop.
|
|
while True:
|
|
with print_lock:
|
|
for thread_id in range(len(thread_status)):
|
|
sys.stdout.write('\033[{0};{1}H{2}'.format(thread_id + 1, info_offset, thread_status[thread_id]))
|
|
sys.stdout.flush()
|
|
time.sleep(0.1)
|
|
|
|
def work_thread(thread_id, options, q):
|
|
global longest_algo
|
|
global info_offset
|
|
global info_offset_found
|
|
global thread_status
|
|
global result_line
|
|
global max_result_line
|
|
global longest_result_line
|
|
|
|
# Save some useful stuff for later.
|
|
#exe_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'bruteforce')
|
|
exe_path = '/home/richard/bruteforce'
|
|
process_args = [exe_path, '[magic]']
|
|
output_base = os.path.join('bruteforce_out', os.path.basename(options['file_path']))
|
|
|
|
while True:
|
|
# Start process.
|
|
proc = subprocess.Popen(process_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
|
|
|
|
# Initialize thread status.
|
|
thread_status[thread_id] = 'init'
|
|
|
|
# Read file into buffer 0.
|
|
proc.stdin.write('0;0\n{0}\n'.format(options['file_path']).encode('utf8', 'ignore'))
|
|
proc.stdin.flush()
|
|
|
|
# Read size of buffer 0.
|
|
file_size = int(proc.stdout.readline().strip())
|
|
if info_offset_found <= 0:
|
|
info_offset_found = info_offset + longest_algo + (len(hex(file_size - 1)) - 2) + 3
|
|
|
|
# Initialize buffer 1.
|
|
proc.stdin.write('1;1;{0}\n'.format(options['decomp_buf_size']).encode('utf8', 'ignore'))
|
|
proc.stdin.flush()
|
|
|
|
# Read size of buffer 1.
|
|
buf_size = int(proc.stdout.readline().strip())
|
|
if buf_size < options['decomp_size']:
|
|
options['decomp_size'] = buf_size
|
|
|
|
# Work until the process dies.
|
|
while True:
|
|
# Receive work from the queue.
|
|
item = q.get()
|
|
if item == None: # special item to stop the loop
|
|
proc.stdin.close()
|
|
return
|
|
|
|
offset, algo = item
|
|
|
|
# Request decompression.
|
|
proc.stdin.write('3;0;{0};{1};1;0;{2};{3}\n'.format(offset, file_size - offset, options['decomp_size'], algo).encode('utf8', 'ignore'))
|
|
proc.stdin.flush()
|
|
|
|
# Set thread status.
|
|
thread_status[thread_id] = algos[algo].ljust(longest_algo) + ' ' + hex(offset)[2:]
|
|
|
|
# Read result.
|
|
try:
|
|
decompressed = proc.stdout.readline()
|
|
except:
|
|
decompressed = b''
|
|
if decompressed:
|
|
decompressed = int(decompressed.strip())
|
|
else:
|
|
try:
|
|
proc.kill()
|
|
except:
|
|
pass
|
|
break
|
|
|
|
# Check for success.
|
|
found = None
|
|
if options['byte_search']:
|
|
# Request byte search.
|
|
proc.stdin.write('4;1;0;{0};{1}\n'.format(options['check_length'], len(options['byte_search'])).encode('utf8', 'ignore'))
|
|
proc.stdin.write(options['byte_search'])
|
|
proc.stdin.flush()
|
|
|
|
# Read result.
|
|
found_offset = int(proc.stdout.readline().strip())
|
|
|
|
# Print result.
|
|
if found_offset > -1:
|
|
found = '{0}+{1}'.format(hex(offset)[2:], hex(found_offset)[2:])
|
|
elif decompressed >= options['check_length']:
|
|
found = hex(offset)[2:]
|
|
|
|
if found:
|
|
# Request buffer read.
|
|
if decompressed > 0:
|
|
request_size = decompressed
|
|
else:
|
|
request_size = buf_size
|
|
proc.stdin.write('2;1;0;{0}\n'.format(request_size).encode('utf8', 'ignore'))
|
|
proc.stdin.flush()
|
|
|
|
# Read data length.
|
|
data_length = int(proc.stdout.readline().strip())
|
|
|
|
# Write buffer data to file.
|
|
f = open(output_base + '.' + found, 'wb')
|
|
f.write(proc.stdout.read(data_length))
|
|
f.close()
|
|
|
|
with print_lock:
|
|
s = algos[algo] + ' ' + found
|
|
if len(s) > longest_result_line:
|
|
longest_result_line = len(s)
|
|
if info_offset_found + len(s) < term_size.columns:
|
|
sys.stdout.write('\033[{0};{1}H{2}'.format(result_line, info_offset_found, s))
|
|
sys.stdout.flush()
|
|
result_line += 1
|
|
if result_line >= term_size.lines:
|
|
result_line = 1
|
|
info_offset_found += longest_result_line + 2
|
|
elif result_line > max_result_line:
|
|
max_result_line = result_line
|
|
|
|
def main():
|
|
global longest_algo
|
|
global info_offset
|
|
global thread_status
|
|
|
|
# Set default options.
|
|
options = {
|
|
'byte_search': None,
|
|
'check_length': 0,
|
|
'decomp_size': 0,
|
|
'decomp_buf_size': 0,
|
|
'threads': 0
|
|
}
|
|
|
|
# Parse arguments.
|
|
args, remainder = getopt.gnu_getopt(sys.argv[1:], 'b:l:d:s:t:', ['byte-search', 'byte-search-length', 'decomp-size', 'decomp-buf-size', 'threads'])
|
|
for opt, arg in args:
|
|
if opt in ('-b', '--byte-search'):
|
|
try:
|
|
options['byte_search'] = eval('b\'' + arg + '\'')
|
|
except:
|
|
pass
|
|
elif opt in ('-l', '--check-length'):
|
|
try:
|
|
options['check_length'] = int(arg)
|
|
except:
|
|
pass
|
|
elif opt in ('-d', '--decomp-size'):
|
|
try:
|
|
options['decomp_size'] = int(arg)
|
|
except:
|
|
pass
|
|
elif opt in ('-s', '--decomp-buf-size'):
|
|
try:
|
|
options['decomp_buf_size'] = int(arg)
|
|
except:
|
|
pass
|
|
elif opt in ('-t', '--threads'):
|
|
try:
|
|
options['threads'] = int(arg)
|
|
except:
|
|
pass
|
|
|
|
if len(remainder) < 1:
|
|
print('Usage: python3 bruteforce.py [-b byte_search_escaped] [-l check_length] [-d decomp_size] [-s decomp_buf_size] [-t threads] file_path')
|
|
return 1
|
|
options['file_path'] = remainder[0]
|
|
|
|
# Set more default options.
|
|
if options['decomp_buf_size'] <= 0:
|
|
options['decomp_buf_size'] = 16777216
|
|
if options['decomp_size'] <= 0 or options['decomp_size'] > options['decomp_buf_size']:
|
|
options['decomp_size'] = options['decomp_buf_size']
|
|
if options['check_length'] <= 0:
|
|
options['check_length'] = options['decomp_size']
|
|
if options['threads'] <= 0:
|
|
options['threads'] = os.cpu_count() or 4
|
|
|
|
# Clear screen.
|
|
print_('\033c')
|
|
|
|
# Create output directory.
|
|
try:
|
|
os.makedirs('bruteforce_out')
|
|
except:
|
|
pass
|
|
|
|
# Prepare queue.
|
|
q = queue.Queue(maxsize=options['threads'])
|
|
|
|
# Start work threads.
|
|
threads = []
|
|
for thread_id in range(options['threads']):
|
|
thread_status.append('')
|
|
thread = threading.Thread(target=work_thread, args=(thread_id, options, q))
|
|
thread.daemon = False
|
|
thread.start()
|
|
threads.append(thread)
|
|
|
|
# Determine algorithms to use.
|
|
for algo in range(len(algos)):
|
|
if not options['byte_search'] and algo != 0: # only lh5 is complex enough to have return values
|
|
algos[algo] = None
|
|
|
|
# Start print thread.
|
|
longest_algo = max(len(algo) for algo in algos if algo)
|
|
info_offset = len(str(options['threads'])) + 2
|
|
thread = threading.Thread(target=print_thread)
|
|
thread.daemon = True
|
|
thread.start()
|
|
|
|
try:
|
|
# Add every offset-algorithm combination.
|
|
file_size = os.path.getsize(options['file_path'])
|
|
for offset in range(file_size):
|
|
for algo in range(len(algos)):
|
|
if algos[algo]:
|
|
q.put((offset, algo))
|
|
|
|
# Stop threads.
|
|
for _ in range(options['threads']):
|
|
q.put(None)
|
|
|
|
# Wait for threads.
|
|
for thread in threads:
|
|
thread.join()
|
|
except KeyboardInterrupt:
|
|
pass
|
|
|
|
# Stop threads again.
|
|
for _ in range(options['threads']):
|
|
q.put(None)
|
|
|
|
# Move cursor to a sane position.
|
|
global result_line
|
|
print_('\033[{0};1H'.format(max(options['threads'], max_result_line) + 1))
|
|
|
|
return 0
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(main())
|