ID2T-toolkit/code/Test/test_FTPWinaXeExploit.py

import os
import random
import unittest
import unittest.mock as mock

from ID2TLib.Controller import *
from Test.Lib import test_pcap, get_sha256


sha_one_attacker_ftp = 'd43ebfd5b0f3c9b6556e6a8d2ae4da4d334ed673881d67720eda16faab1b498e'



class UnitTestFTPWinaXeExploit(unittest.TestCase):

    def clean_up(self, controller):
        os.remove(controller.pcap_dest_path)
        os.remove(controller.label_manager.label_file_path)


    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop')
    @mock.patch('ID2TLib.Utility.get_rnd_bytes')
    def test_one_attacker_ftp(self,mock_rnd_x86_nop,mock_rnd_bytes):
        random.seed(5)
        mock_rnd_x86_nop.return_value = b''
        mock_rnd_bytes.return_value = b''
        controller = Controller(pcap_file_path=test_pcap, do_extra_tests=False)
        controller.load_pcap_statistics(False, False, False)
        controller.process_attacks([['FTPWinaXeExploit' ]])
        self.assertEqual(get_sha256(controller.pcap_dest_path), sha_one_attacker_ftp)
        self.clean_up(controller)


if __name__ == '__main__':
    unittest.main()