malware_analysis_toolkit.py

import argparse
import logging
import os
import pefile
import frida

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def static_analysis(file_path):
    pe = pefile.PE(file_path)

    logger.info("---------- Static Analysis ----------")
    logger.info("File: %s", file_path)
    logger.info("ImageBase: %s", hex(pe.OPTIONAL_HEADER.ImageBase))
    logger.info("Entry Point: %s", hex(pe.OPTIONAL_HEADER.AddressOfEntryPoint))
    logger.info("Number of Sections: %s", pe.FILE_HEADER.NumberOfSections)
    logger.info("Sections:")
    for section in pe.sections:
        logger.info("- %s", section.Name.decode().rstrip('\x00'))

    # Extract and print imported functions
    logger.info("\n---------- Imported Functions ----------")
    if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            logger.info("%s:", entry.dll.decode())
            for imp in entry.imports:
                if imp.name:
                    logger.info("- %s", imp.name.decode())
                else:
                    logger.info("- %s", hex(imp.address))
    else:
        logger.info("No imported functions found.")


def dynamic_analysis(file_path):
    def on_message(message, data):
        if message["type"] == "send":
            logger.info("[+] Message from script: %s", message["payload"])
        elif message["type"] == "error":
            logger.error("[-] Error from script: %s", message["description"])

    with open(os.path.join(os.path.dirname(__file__), "dynamic_analysis.js"), "r") as f:
        script_code = f.read()

    process = frida.attach(file_path)
    script = process.create_script(script_code)
    script.on("message", on_message)
    script.load()
    process.resume()
    script.join()


def extract_strings(file_path):
    with open(file_path, "rb") as f:
        content = f.read()
        printable_strings = [s.decode(errors="ignore") for s in content.split(b"\x00") if s.strip()]
        if printable_strings:
            logger.info("\n---------- Extracted Strings ----------")
            for string in printable_strings:
                logger.info(string)


def malware_analysis(file_path):
    try:
        file_path = os.path.abspath(file_path)
        if not os.path.isfile(file_path):
            raise FileNotFoundError(f"Error: File not found at {file_path}")

        if not pefile.is_pefile(file_path):
            raise ValueError("Error: Not a valid PE file")

        static_analysis(file_path)
        dynamic_analysis(file_path)
        extract_strings(file_path)

    except FileNotFoundError as e:
        logger.error(e)
    except pefile.PEFormatError as e:
        logger.error("Error: Invalid PE file format")
    except ValueError as e:
        logger.error(e)
    except Exception as e:
        logger.error("An error occurred: %s", str(e))


def add_arguments(parser):
    parser.add_argument("file_path", help="Path of the file to analyze")


def main():
    parser = argparse.ArgumentParser(description="Malware Analysis Tool")
    add_arguments(parser)
    args = parser.parse_args()

    malware_analysis(args.file_path)


if __name__ == "__main__":
    main()