StaticMalwareAnalysis.py

import os
import hashlib
from datetime import datetime
import shutil
from config import *
from utils import *
import pdfkit
import pefile
import requests
import bs4

class StaticMalwareAnalysis:
    def __init__(self, sample_name):
        self.path = config['malware-sample-folder']
        self.saved_analysis_folder = config['saved-analysis']
        self.sample_file = sample_name                                  # malware-folder/mal.exe
        self.sample_name = os.path.basename(sample_name)                # mal.exe
        self.sample_name_dir = ""                                       # saved-analysis/mal.exe-timestamp/
        self.information = {
            "archirecture": "",
            "sha256_hash" : "",
            "malicious_api": [],
            "other_api" : [],
            "virustotal_results": None,
            "strings_file" : "",
            "sections": [],
            "saved_to": ""
        }
        
    @staticmethod
    def check_malware_sample_exists():
        path = config['malware-sample-folder']
        if not os.path.exists(path):
            os.mkdir(path)
            print("Malware Sample Folder doesnot exist. Created folder.")
            return False
        elif len(os.listdir(path)) == 0:
            print("Malware Sample directory is empty. Please add malware to start static analysis")
            return False
        return True
    
    
    def create_directory(self):
        if not os.path.exists(self.saved_analysis_folder):
            os.mkdir(self.saved_analysis_folder)
        sample_name_dir = self.sample_name + "_" + str(datetime.now()).replace(" ","-").replace(":","").replace(".","")
        os.mkdir(os.path.join(self.saved_analysis_folder,sample_name_dir))
        self.sample_name_dir = os.path.join(self.saved_analysis_folder,sample_name_dir)
        return sample_name_dir
    
    def get_sha256sum(self):
        hash_function = hashlib.sha256()
        with open(self.sample_file, "rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                hash_function.update(byte_block)
            hash_value = hash_function.hexdigest()
        self.information['sha256_hash'] = hash_value
        return hash_value
    
    def extract_strings(self):
        output_file_txt = os.path.join(self.sample_name_dir,f'{self.sample_name}-strings.txt')
        output_file_html = os.path.join(self.sample_name_dir,f'{self.sample_name}-strings.html')
        command = f"flarestrings -n 6 {self.sample_file} | rank_strings --scores > {output_file_txt}"
        os.system(command)

        strings = open(output_file_txt,"r")
        with open(output_file_html, "w") as html_file:
            for lines in strings.readlines():
                html_file.write("<pre>" + escape_html(lines) + "</pre>\n")
        self.information['strings_file'] = f'{self.sample_name}-strings.html'
        print(f"Strings extracted to file : {output_file_html}")

    def malAPICheck(self,api):
        api_request = requests.get("https://malapi.io/winapi/" + api)
        api_request.raise_for_status()
        api_data = bs4.BeautifulSoup(api_request.text, 'html.parser')
        
        if "404 Not Found" in api_data.getText():
            print(f"API NOT FOUND - {api}")
            return "No result", "No result"

        details = api_data.select('.detail-container .content')
        api_description = details[1].getText().lstrip().rstrip()
        api_associated_attacks = " ".join(details[3].getText().lstrip().rstrip().split())
        print(f"API FOUND - {api}")
        return api_description, api_associated_attacks
    
    def extractPEINFO(self):
        try:
            pe = pefile.PE(self.sample_file)

            if hex(pe.FILE_HEADER.Machine) == '0x14c':
                self.information['archirecture'] = "32 bit binary"
            else:
                self.information['archirecture'] = "64 bit binary"

            for section in pe.sections:
                section_name = section.Name.decode().rstrip('\x00')
                self.information['sections'].append((
                    f"{section_name}",
                    hex(section.Misc_VirtualSize),
                    hex(section.SizeOfRawData),
                    str(hex(section.Misc_VirtualSize) > hex(section.SizeOfRawData))
                ))

            pe.parse_data_directories()
            for entry in pe.DIRECTORY_ENTRY_IMPORT:
                for imp in entry.imports:
                    api = imp.name.decode("utf-8")
                    try:
                        if api[0] != '_':
                            desc, tags = self.malAPICheck(api)
                            if desc == "No result" or tags == "No result":
                                self.information['other_api'].append(api)
                            else:
                                self.information['malicious_api'].append((api, desc.replace(",",""), tags))
                    except:
                        self.information['other_api'].append(api)
            return self.information
        except Exception as e:
            print("Error occurred in collecting PE Info")
            print(e)
            return False

    def move_malware(self):
        final_path = os.path.join(self.sample_name_dir, "Malware." + self.sample_name + ".malz")
        shutil.copyfile(self.sample_file, final_path)
        self.information['saved_to'] = final_path
        try:
            os.unlink(self.sample_file)
        except Exception as e:
            print(f"File delete failed - {self.sample_file}")
        return final_path
    
    def set_virustotal_results(self, results):
        self.information['virustotal_results'] = results
    
    def generate_virustotal_table(self):
        if self.information['virustotal_results'] == None:
            return ""
        else:
            html_table = ""
            magic_header = self.information['virustotal_results'].get("magic_header", None)
            if magic_header != None:
                html_table += f"""
                    <tr>
                        <th> Magic Header </th>
                        <td> {magic_header} </td>
                    </tr>
                """
            
            times_submitted = self.information['virustotal_results'].get("times_submitted", None)
            if times_submitted != None:
                html_table += f"""
                    <tr>
                        <th> Times Submitted </th>
                        <td> {times_submitted} </td>
                    </tr>
                """
            
            packers = self.information['virustotal_results'].get("packers", None)
            if packers != None:
                html_table += f"""
                    <tr>
                        <th> Packer PEiD </th>
                        <td> {packers.get("PEiD","-")} </td>
                    </tr>
                """

            names = self.information['virustotal_results'].get("names", None)
            if names != None and len(names) > 0: 
                html_table += f"""
                    <tr>
                        <th> Names </th>
                        <td> {", ".join(names)} </td>
                    </tr>
                """

            threat_classification = self.information['virustotal_results'].get("threat_classification", None)
            if threat_classification != None : 
                html_table += f"""
                    <tr>
                        <th> Threat Classification sugested label </th>
                        <td> {threat_classification.get("suggested_threat_label","-")} </td>
                    </tr>
                """
            
            yara_rules = self.information['virustotal_results'].get("yara_rules", None)
            if yara_rules != None and len(yara_rules) > 0:
                for i in yara_rules:
                    html_table += """
                        <tr>
                            <th colspan="2" > Yara Rules</th>
                        </tr>
                    """
                    rule_name = i.get('rule_name','-')
                    rule_source = i.get('source','-')
                    html_table += f"""
                        <tr>
                            <th> {rule_name} </th>
                            <td> {rule_source} </td>
                        </tr>
                    """

            if html_table == "":
                return ""
            else:
                return  f"""
                    <br><hr/>
                    <h3> Virus Total Results </h3>
                    <table>
                        <tbody>
                            {html_table}
                        <tbody>
                    </table>
                """

    def create_pdf(self):
        malicious_api_tbody = ""
        malicious_api = self.information['malicious_api']
        for row in malicious_api:
            name = f'<a href="https://malapi.io/winapi/{row[0]}" target="_blank" rel="noopener noreferrer">{row[0]}</a>'
            
            malicious_api_tbody += f"""
                <tr>
                    <td>{ name }</td>
                    <td>{row[1]}</td>
                    <td>{row[2]}</td>
                </tr>
            """

        other_api_body = ", ".join(self.information['other_api'])

        sections_tbody = ""
        sections = self.information['sections']
        for row in sections:
            sections_tbody += f"""
                <tr>
                    <td>{row[0]}</td>
                    <td>{row[1]}</td>
                    <td>{row[2]}</td>
                    <td>{row[3]}</td>
                </tr>
            """

        styles = """
            table, th, td { border: 1px solid black;border-collapse: collapse;}
            table { margin-left: auto; margin-right: auto; }
            tr,h3 { text-align: center; }
            th { padding : 10px; }
            td { padding : 3px; }
        """

        html = f"""
            <html lang="en">
                <head>
                    <meta charset="UTF-8">
                    <meta http-equiv="X-UA-Compatible" content="IE=edge">
                    <meta name="viewport" content="width=device-width, initial-scale=1.0">
                    <title>{self.sample_name}</title>
                    <style> {styles} </style>
                </head>
                <body>
                    <h3> General Information </h3>
                    <table>
                        <thead>
                            <tr>
                                <th> SHA 256 HASH </th>
                                <th> {f'<a href="https://www.virustotal.com/gui/search/{self.information["sha256_hash"]}" target="_blank" rel="noopener noreferrer">{self.information["sha256_hash"]}</a>'} </th>
                            </tr>
                            <tr>
                                <th> Architecture </th>
                                <th> {self.information['archirecture']} </th>
                            </tr>
                            <tr>
                                <th colspan="2" > {f'<a href="{self.information["strings_file"]}" target="_blank" rel="noopener noreferrer"> Strings </a>'} </th>
                            </tr>
                        </thead>
                    </table>
                    <br><hr/>
                    <h3> Malicious API's </h3>
                    <table>
                        <thead>
                            <tr>
                                <th>Name</th>
                                <th>Tags</th>
                                <th>Malicious</th>
                            </tr>
                        </thead>
                        <tbody>
                            {malicious_api_tbody}
                        </tbody>
                    </table>
                    <br><hr/>
                    <h3> Other API's </h3>
                    <p> {other_api_body} </p>
                    <br><hr/>
                    <h3> Sections </h3>
                    <table>
                        <thead>
                            <tr>
                                <th>Name</th>
                                <th>Virtual Size</th>
                                <th>Raw Data Size</th>
                                <th>Virtual Size >  Raw Data</th>
                            </tr>
                        </thead>
                        <tbody>
                            {sections_tbody}
                        </tbody>
                    </table>
                    {self.generate_virustotal_table()}
                </body>
            </html>
        """

        path_to_file = os.path.join(self.sample_name_dir,f'{self.sample_name}.html')
        with open(path_to_file,"w") as html_file:
            html_file.write(html)
        
        path_to_wkhtmltopdf = config['path_to_wkhtmltopdf']
        if not os.path.exists(path_to_wkhtmltopdf):
            print("wkhtmltopdf binary not found, skipping PDF creation")
            print(f"HTML file saved in {path_to_file}")
        else:
            pdf_config = pdfkit.configuration(wkhtmltopdf=path_to_wkhtmltopdf)
            pdfkit.from_file(
                path_to_file, 
                output_path=os.path.join(self.sample_name_dir,f'{self.sample_name}.pdf'), 
                configuration=pdf_config
            )