Import_Table_catalog

PE_parse_and_catalog.py

import sqlite3
import argparse
import os
import subprocess
import re

# This script iterates over all executable files within given directory (e.g. C:\Windows, C:\Program Files) 
# and runs `rabin2 -i EXEPATH` in order to obtain its import table, then saves the import table into the database (SQLite).
# We will use this database later for further analysis purposes.

if __name__ == "__main__":
	database = r"import_tables.db"
	parser = argparse.ArgumentParser()
	group = parser.add_mutually_exclusive_group(required=True)
	group.add_argument('--create','-c', help='Create database', action='store_true')
	group.add_argument('--import','-i', help='Import executable', action='store_true')
	group.add_argument('--list','-l', help='List all executables', action='store_true')
	parser.add_argument('--function_name','-F') # for get_executables
	parser.add_argument('--path','-P') # for import
	args = parser.parse_args()

def create_connection(db_file):
	conn = None
	try:
		conn = sqlite3.connect(db_file)
	except Error as e:
		print(e)
	return conn


def create_db(conn):
	createContentTable="""CREATE TABLE IF NOT EXISTS PE_import (
id integer PRIMARY KEY,
image_path text NOT NULL,
library_name text NOT NULL,
function_name text NOT NULL);"""
	try:
		c = conn.cursor()
		c.execute(createContentTable)
	except Error as e:
		print(e)

def insert_pe(conn, image_path, library_name, function_name):  # DONE
	rowid = 0
	try:
		c = conn.cursor()
		c.execute("INSERT INTO PE_import (image_path, library_name, function_name) VALUES ('%s','%s','%s')" % (image_path, library_name, function_name))
		rowid = c.lastrowid
	except Error as e:
		print(e)
	return rowid

def list_executables(conn, function_name): # DONE, get all images having the particular function in their import table
	rows = None
	try:
		c = conn.cursor()
		c.execute("SELECT image_path, library_name FROM PE_import WHERE function_name = '%s'" % function_name)
		rows = c.fetchall()
	except Error as e:
		print(e)
	for row in rows:
		print(row[0][0]+"\t("+row[0][1]+")") # image_path + library_name (e.g. C:\Windows\System32\notepad.exe kernel32.dll)

def import_executable(conn, path): # let's see the output, if it's correct, let's parse it and pass to insert_pe
	rabin_import = subprocess.run(["rabin2", "-i", path], stdout=subprocess.PIPE, text=True)
	entries = rabin_import.stdout.splitlines()
	for entry in entries:
		m = re.search('FUNC\s+(\S+)\s+(\S+)', entry)
		if(m):
			DLL = m.group(1)
			function_name = m.group(2)
			print("%s\t%s" % (DLL, function_name))
			insert_pe(conn, path, DLL, function_name)
	conn.commit()
	return

conn = create_connection(database)

if (args.create):
	print("[+] Creating database")
	create_db(conn)
elif (args.list):
	if(args.function_name is None):
		parser.error("--list requires --function_name")
	else:
		print("[+] The following executables utilize the %s function:")
		list_executables(conn, args.function_name)
elif ('import' in vars(args).keys()):
	if(args.path is None):
		parser.error("--import requires --path.")
	else:
		#print("[+] Importing the Import Table entries from %s " % args.path)
		import_executable(conn, args.path)