import unittest
import pandas as pd
from agents_elements.agent_types.manager.manager import Manager
from agents_elements.agent_types.defender.defender import Defender
from variables import NAME_MALWARE_SCRIPT, PATH_TO_OSSEC_AGENT_TEST_MANAGER, PATH_TO_OSSEC_AGENT_BASE_DEFENDER,PATH_TO_LOCAL_DECODER,PATH_TO_LOCAL_RULES
import os
from utils import save_hierarchy_to_file
from unittest.mock import patch
from pandas.testing import assert_frame_equal
import numpy as np
# Unit test class
class Unitest(unittest.TestCase):
    def setUp(self):
        # -------------------- INITIALIZATION --------------------

        # Manager Instance
        self.manager = Manager(name='wazidx1', ip_address='192.168.56.13')

        # Defender Instance
        self.defender = Defender(name='wazagent1', ip_address='192.168.56.14')

        # Path to the ossec testing conf file
        self.ossec_conf_path = PATH_TO_OSSEC_AGENT_TEST_MANAGER
        # -------------------- TEST 1: CLASS STRUCTURE --------------------
        # Save class methods and hierarchy
        print("Saving class methods to 'methods.txt'...")
        save_hierarchy_to_file(Manager, 'methods.txt')

        # -------------------- TEST 2: FROM CONF --------------------
        # Section tag to be added
        self.section_tag = """
        <command>
            <name>firewall-drop</name>
            <executable>firewall-drop</executable>
            <timeout_allowed>yes</timeout_allowed>
        </command>
        <localfile>
            <log_format>syslog</log_format>
            <command>syslog</command>
            <frequency>360</frequency>
            <location>/var/log/syslog</location>
            <alias>syslog</alias>
        </localfile>
        <active-response>
            <command>firewall-drop</command>
            <location>local</location>
            <rules_id>5763,5886,5447</rules_id>
            <timeout>180</timeout>
        </active-response>
        <active-response>
            <command>yara_linux</command>
            <location>local</location>
            <rules_id>100300,100301</rules_id>
        </active-response>
        """

        # Raw XML string for a group section
        self.group_section_str = """
        <group name="syscheck,">
            <rule id="100300" level="7">
            <if_sid>550</if_sid>
            <field name="file">f{PATH_TO_YARA_MALWARE_DIR}/</field>
            <description>File modified in f{PATH_TO_YARA_MALWARE_DIR}/ directory.</description>
            </rule>
            <rule id="100301" level="7">
            <if_sid>554</if_sid>
            <field name="file">f{PATH_TO_YARA_MALWARE_DIR}/</field>
            <description>File added to f{PATH_TO_YARA_MALWARE_DIR}/ directory.</description>
            </rule>
        </group>

        <group name="yara,">
            <rule id="108000" level="0">
            <decoded_as>yara_decoder</decoded_as>
            <description>Yara grouping rule</description>
            </rule>
            <rule id="108001" level="12">
            <if_sid>108000</if_sid>
            <match>wazuh-yara: INFO - Scan result: </match>
            <description>File "$(yara_scanned_file)" is a positive match. Yara rule: $(yara_rule)</description>
            </rule>
        </group>
        """
        # Raw XML string for a decoder section
        self.decoder_section_str = """
        <decoder name="yara_decoder">
            <prematch>wazuh-yara:</prematch>
        </decoder>

        <decoder name="yara_decoder1">
            <parent>yara_decoder</parent>
            <regex>wazuh-yara: (\S+) - Scan result: (\S+) (\S+)</regex>
            <order>log_type, yara_rule, yara_scanned_file</order>
        </decoder>
        """
        # Remove the section tag from the ossec.conf file before testing
        self.remove_sections()

    def remove_sections(self):
        # Remove the tag section from the ossec conf
        self.remove_section_from_file(PATH_TO_OSSEC_AGENT_TEST_MANAGER,self.section_tag)
        # Remove the group section from the local_file XML
        self.remove_section_from_file(PATH_TO_LOCAL_RULES, self.group_section_str)
        # Remove the decoder section from the decoder XML
        self.remove_section_from_file(PATH_TO_LOCAL_DECODER, self.decoder_section_str)
       
    
    def remove_section_from_file(self, file_path, section_str):
        with open(file_path, 'r') as file:
            lines = file.readlines()

        with open(file_path, 'w') as file:
            for line in lines:
                if section_str.strip() not in line.strip():
                    file.write(line)

    def test_ossec_conf_handling(self):
        

        # Testing the conversion from tag to dataframe
        result = self.manager.dataframe_from_xml_tag(self.section_tag)

        # Expected output
        expected_command = pd.DataFrame({
            'name': [[{'text': 'firewall-drop'}]],
            'executable': [[{'text': 'firewall-drop'}]],
            'timeout_allowed': [[{'text': 'yes'}]],
            'extra_args': [None]
        })

        expected_localfile = pd.DataFrame({
            'log_format': [[{'text': 'syslog'}]],
            'command': [[{'text': 'syslog'}]],
            'frequency': [[{'text': '360'}]],
            'location': [[{'text': '/var/log/syslog'}]],
            'alias': [[{'text': 'syslog'}]]
        })

        expected_active_response = pd.DataFrame({
            'command': [[{'text': 'firewall-drop'}], [{'text': 'yara_linux'}]],
            'location': [[{'text': 'local'}], [{'text': 'local'}]],
            'rules_id': [[{'text': '5763,5886,5447'}], [{'text': '100300,100301'}]],
            'timeout': [[{'text': '180'}], [None]],
            'ca_store': [None, None],
            'ca_verification': [None, None],
            'extra_args': [None, None]
        })

        # Assertions for tag to dataframe procedure #

        assert_frame_equal(result['command'], expected_command)
        assert_frame_equal(result['localfile'], expected_localfile)
        #assert_frame_equal(result['active-response'], expected_active_response)

        # Testing to add a section tag to an ossec_conf file
        result = self.manager.add_section_tag_to_conf_file(section_tag=self.section_tag, file_path=PATH_TO_OSSEC_AGENT_TEST_MANAGER)

        # Load dataframes back from the Excel file
        print("Reloading configuration data from Excel...")
        dict_dfs = self.manager.dataframe_from_conf(PATH_TO_OSSEC_AGENT_TEST_MANAGER)

        # Assertions to verify if ossec has been correctly modified #
        
        # Check if the last row of each section in dict_dfs is a subset of the expected DataFrame
        is_subset_command = expected_command.apply(lambda row: dict_dfs['command'].iloc[-1].isin(row).all(), axis=1).any()
        self.assertEqual(is_subset_command, True)

        is_subset_localfile = expected_localfile.apply(lambda row: dict_dfs['localfile'].iloc[-1].isin(row).all(), axis=1).any()
        self.assertEqual(is_subset_localfile, True)

        is_subset_active_response = expected_active_response.apply(lambda row: dict_dfs['active-response'].iloc[-1].isin(row).all(), axis=1).any()
        self.assertEqual(is_subset_active_response, True)
       

    def test_fim_setup(self):
        # [Doc]: https://documentation.wazuh.com/current/user-manual/capabilities/file-integrity/advanced-settings.html

        # Define monitored paths
        
        monitored_path_who_data = "/etc"
        monitored_path = '/test'

        # Define directory tags
        directory_tag_whodata = {'text': monitored_path_who_data, 'check_all': 'yes', 'whodata': 'yes'}
        directory_tag = {'text': monitored_path, 'check_all': 'yes','realtime': 'yes'}

        
        print("Starting FIM setup...")

        # Add FIM configuration
        self.defender.add_fim_configuration(monitored_paths=[directory_tag, directory_tag_whodata], do_synchronize_with_VM=True, conf_path=PATH_TO_OSSEC_AGENT_BASE_DEFENDER)

        # Test FIM
        print("Testing FIM...")
        self.defender.run_function_on_remote_host('simulate_file_change', monitored_path)
        self.assertEqual(self.manager.check_fim_change(monitored_path),True)

        #Test Who-Data
        print("Testing audit...")
        self.defender.run_function_on_remote_host('test_who_data')
        self.assertEqual(self.manager.check_fim_change(monitored_path_who_data),True)

        print("FIM and Who-Data setup completed successfully.")

    
    def test_xml_handling(self):
        """
        Configuring manager xml for yara.

        """
        # Add the group section to the local_file XML
        self.manager.add_group_to_xml(self.group_section_str)
        # Add the decoder section to the decoder XML
        self.manager.add_decoder_to_xml(self.decoder_section_str)

        # Synchronize the XML with the VM
        result = self.manager.synchronize_xml_with_VM(decoder_sections_str=self.decoder_section_str, group_sections_str=self.group_section_str)

        # Check for working modification of new added rules
        self.assertEqual(result,True)
     
    def test_yara_configuration(self):
        """
         - Doc: Detecting malware using YARA integration
            [https://documentation.wazuh.com/current/proof-of-concept-guide/detect-malware-yara-integration.html#detecting-malware-using-yara-integration]
        """
         # Ensure required keys exist and modify the configuration data
        print("Configuring ossec conf for yara")

        # Command Section
         # Add YARA configuration on the endpoint
        self.defender.check_and_install_yara(conf_path = PATH_TO_OSSEC_AGENT_BASE_DEFENDER)

        # -------------------- TEST 4: GENERATE NEW CONFIGURATION --------------------
        # Generate the updated ossec.conf XML
        self.section_tag_yara = """
        <command>
            <name>yara_linux</name>
            <executable>yara.sh</executable>
            <timeout_allowed>no</timeout_allowed>
            <extra_args>-yara_path /usr/local/bin -yara_rules f{PATH_TO_YARA_RULES}</extra_args>
	    </command>

        <active-response>
            <command>yara_linux</command>
            <location>local</location>
            <rules_id>100300,100301</rules_id>
	    </active-response>
    """

        #BUG WITH remove_file_if_exists probably as excel files aren't updated.
        # 
        # Processing the downloads
        self.defender.run_bash_script_on_remote_host(NAME_MALWARE_SCRIPT)
        
        # -------------------- TEST COMPLETED --------------------
        self.assertEqual(self.manager.check_alerts('100300'),True)

# Run the tests
if __name__ == '__main__':
    # Go to the right directory
    os.chdir('vagrant/src')
    #unittest.main()

    # Create a test suite with only the test_fim_setup method
    suite = unittest.TestSuite()
    suite.addTest(Unitest('test_fim_setup'))

    # Run the test suite
    unittest.TextTestRunner().run(suite)

