import unittest
import unittest.mock as mock

from Test.GenericTest import GenericTest
from Test.Lib import *

sha_ftp_basic = '941947ccc42ea10e724d2a20626882130d62fc5dbbe007095a90f67a943ab3bf'
sha_ftp_most_used_ip = '941947ccc42ea10e724d2a20626882130d62fc5dbbe007095a90f67a943ab3bf'
sha_ftp_mac = 'c2e83e62bb8a15402725faef47a53c6e5afa3dd82a17435d48000058976160cb'
sha_ftp_random_ip_src = '41ae677b553064428905682f6a17447850cc4c1b617c337e046ee6e50f51217b'
sha_not_empty_custom_payload_empty_file = '369d59174de5f01787ea623673f320e8342ddd6be9761edb607bf635f44a3749'
sha_empty_custom_payload_not_empty_file = '9d3ec2451b05acc72b99b40309b714bc015b6d12b5477f6490cd2f9ba8f1ffa8'
sha_valid_ip = '941947ccc42ea10e724d2a20626882130d62fc5dbbe007095a90f67a943ab3bf'

"""
Name                             Stmts   Miss  Cover   Missing
--------------------------------------------------------------------------------------------
Attack/FTPWinaXeExploit.py         141     14    99%    67
"""


class UnitTestFTPWinaXeExploit(GenericTest):

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=get_x86_nop)
    def test_ftp_basic(self, mock_get_rnd_x86_nop, mock_get_rnd_bytes):
        self.generic_test([['FTPWinaXeExploit']], sha_ftp_basic)

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=get_x86_nop)
    @mock.patch('ID2TLib.Statistics.Statistics.get_most_used_ip_address')
    def test_ftp_most_used_ips(self,mock_most_used_ip_address, mock_get_rnd_x86_nop, mock_get_rnd_bytes):
        mock_most_used_ip_address.return_value = test_pcap_ips
        self.generic_test([['FTPWinaXeExploit']], sha_ftp_most_used_ip)

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=get_x86_nop)
    @mock.patch('ID2TLib.Statistics.Statistics.get_mac_address')
    def test_ftp_mac(self, mock_mac_address, mock_get_rnd_x86_nop, mock_get_rnd_bytes):
        mock_mac_address.return_value = test_pcap_empty
        self.generic_test([['FTPWinaXeExploit']], sha_ftp_mac)

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=get_x86_nop)
    def test_ftp_random_ip_src(self, mock_get_rnd_x86_nop, mock_get_rnd_bytes):
        self.generic_test([['FTPWinaXeExploit', 'ip.src.shuffle=1']], sha_ftp_random_ip_src)

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=get_x86_nop)
    def test_ftp_not_empty_custom_payload_empty_file(self, mock_get_rnd_x86_nop, mock_get_rnd_bytes):
        self.generic_test([['FTPWinaXeExploit', 'custom.payload=1']], sha_not_empty_custom_payload_empty_file)

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=get_x86_nop)
    @mock.patch('ID2TLib.Utility.check_payload_len')
    @mock.patch('ID2TLib.Utility.get_bytes_from_file', return_value=b'AAAAA')
    def test_ftp_empty_custom_payload_not_empty_file(self, mock_bytes_from_file, mock_payload_len, mock_get_rnd_x86_nop,
                                                     mock_get_rnd_bytes):
        self.generic_test([['FTPWinaXeExploit', 'custom.payload.file=1']], sha_empty_custom_payload_not_empty_file)

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=get_x86_nop)
    @mock.patch('Attack.BaseAttack.BaseAttack.is_valid_ip_address', return_values=[False, True])
    def test_ftp_invalid_ip(self,mock_valid_ip_check, mock_get_rnd_x86_nop, mock_get_rnd_bytes):
        self.generic_test([['FTPWinaXeExploit']], sha_valid_ip)


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