mirror of
https://github.com/dobin/SuperMega
synced 2026-06-02 17:27:10 +00:00
refactor: bix cleanup
This commit is contained in:
@@ -1,84 +1,18 @@
|
||||
import subprocess
|
||||
import os
|
||||
import time
|
||||
import shutil
|
||||
import pathlib
|
||||
import sys
|
||||
import pefile
|
||||
import glob
|
||||
import logging
|
||||
|
||||
from config import config
|
||||
from project import project
|
||||
from pehelper import *
|
||||
from defs import *
|
||||
|
||||
logger = logging.getLogger("Helper")
|
||||
|
||||
SHC_VERIFY_SLEEP = 0.1
|
||||
|
||||
|
||||
verify_filename = r'C:\Temp\a'
|
||||
build_dir = "build"
|
||||
|
||||
|
||||
def remove_trailing_null_bytes(data):
|
||||
for i in range(len(data) - 1, -1, -1):
|
||||
if data[i] != b'\x00'[0]: # Check for a non-null byte
|
||||
return data[:i + 1]
|
||||
return b'' # If the entire sequence is null bytes
|
||||
|
||||
|
||||
def get_code_section_data(pe_file):
|
||||
try:
|
||||
# Load the PE file
|
||||
pe = pefile.PE(pe_file)
|
||||
|
||||
# Iterate over the sections
|
||||
#for section in pe.sections:
|
||||
# # Check if this is the code section
|
||||
# if '.text' in section.Name.decode().rstrip('\x00'):
|
||||
# data = section.get_data()
|
||||
# data = remove_trailing_null_bytes(data)
|
||||
# logger.info(" > 0x{:X} Code Size: {} (raw code section size: {})".format(
|
||||
# section.VirtualAddress,
|
||||
# len(data), section.SizeOfRawData))
|
||||
# return data
|
||||
|
||||
section = get_code_section(pe)
|
||||
if section == None:
|
||||
raise Exception("Code section not found.")
|
||||
|
||||
logger.info("--[ Code section: {}".format(section.Name.decode().rstrip('\x00')))
|
||||
data = section.get_data()
|
||||
data = remove_trailing_null_bytes(data)
|
||||
logger.info(" > 0x{:X} Code Size: {} (raw code section size: {})".format(
|
||||
section.VirtualAddress,
|
||||
len(data), section.SizeOfRawData))
|
||||
return data
|
||||
|
||||
except FileNotFoundError:
|
||||
logger.info(f"File not found: {pe_file}")
|
||||
except pefile.PEFormatError:
|
||||
logger.info(f"Invalid PE file: {pe_file}")
|
||||
|
||||
|
||||
def write_code_section(pe_file, new_data):
|
||||
# Load the PE file
|
||||
pe = pefile.PE(pe_file)
|
||||
|
||||
# Iterate over the sections
|
||||
for section in pe.sections:
|
||||
# Check if this is the code section
|
||||
if '.text' in section.Name.decode().rstrip('\x00'):
|
||||
file_offset = section.PointerToRawData
|
||||
|
||||
with open(pe_file, 'r+b') as f:
|
||||
f.seek(file_offset)
|
||||
f.write(new_data)
|
||||
#logger.info("Successfully overwritten the .text section with new data.")
|
||||
break
|
||||
|
||||
|
||||
def clean_files():
|
||||
logger.info("--[ Remove old files")
|
||||
|
||||
@@ -89,13 +23,13 @@ def clean_files():
|
||||
"mllink$.lnk",
|
||||
|
||||
# out/ stuff
|
||||
os.path.join(build_dir, "main.asm"),
|
||||
os.path.join(build_dir, "main.bin"),
|
||||
os.path.join(build_dir, "main.c"),
|
||||
os.path.join(build_dir, "peb_lookup.h"),
|
||||
#os.path.join(build_dir, "main.exe"),
|
||||
os.path.join(project.build_dir, "main.asm"),
|
||||
os.path.join(project.build_dir, "main.bin"),
|
||||
os.path.join(project.build_dir, "main.c"),
|
||||
os.path.join(project.build_dir, "peb_lookup.h"),
|
||||
#os.path.join(project.build_dir, "main.exe"),
|
||||
|
||||
verify_filename,
|
||||
project.verify_filename,
|
||||
]
|
||||
for file in files_to_clean:
|
||||
pathlib.Path(file).unlink(missing_ok=True)
|
||||
@@ -133,7 +67,7 @@ def try_start_shellcode(shc_file):
|
||||
subprocess.run([
|
||||
config.get["path_runshc"],
|
||||
shc_file,
|
||||
]) # , check=True
|
||||
])
|
||||
|
||||
|
||||
def file_readall_text(filepath) -> str:
|
||||
@@ -155,4 +89,4 @@ def delete_all_files_in_directory(directory_path):
|
||||
os.remove(file_path)
|
||||
#logger.info(f"Deleted {file_path}")
|
||||
except Exception as e:
|
||||
logger.info(f"Error deleting {file_path}: {e}")
|
||||
logger.info(f"Error deleting {file_path}: {e}")
|
||||
|
||||
Reference in New Issue
Block a user