import unittest
import unittest.mock as mock

import ID2TLib.TestLibrary as Lib
import Test.GenericTest as GenericTest

sha_ftp_basic = 'ad9bc7b55c3b0365c0f02ae9b9b7aafdb43acbdd8c8c274d30cb286821e772cc'
sha_ftp_most_used_ip = 'ad9bc7b55c3b0365c0f02ae9b9b7aafdb43acbdd8c8c274d30cb286821e772cc'
sha_ftp_mac = '388831100c907cfc6815bcc1869f30d937be29091dd8e54a734eb52f14a23f3c'
sha_ftp_random_ip_src = 'b18c0f1d15f1afb239116e1ccec20b03716412eea58ca969f7d2ede1749409e3'
sha_not_empty_custom_payload_empty_file = '41186fc804fb2a8fb3605be3246a5246be927e3187ea82bd2fbe2097643863a8'
sha_empty_custom_payload_not_empty_file = 'b1f43c3147dd3684b1db4d7d370801f25de693b632b97a95b933a4d296094f31'
sha_valid_ip = 'ad9bc7b55c3b0365c0f02ae9b9b7aafdb43acbdd8c8c274d30cb286821e772cc'

"""
Name                             Stmts   Miss  Cover   Missing
--------------------------------------------------------------------------------------------
Attack/FTPWinaXeExploit.py         141     14    99%    67
"""
# TODO: get 100% coverage


class UnitTestFTPWinaXeExploit(GenericTest.GenericTest):

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=Lib.get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=Lib.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=Lib.get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=Lib.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 = Lib.test_pcap_ips
        self.generic_test([['FTPWinaXeExploit']], sha_ftp_most_used_ip)

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=Lib.get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=Lib.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 = Lib.test_pcap_empty
        self.generic_test([['FTPWinaXeExploit']], sha_ftp_mac)

    @mock.patch('ID2TLib.Utility.get_rnd_bytes', side_effect=Lib.get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=Lib.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=Lib.get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=Lib.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=Lib.get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=Lib.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=Lib.get_bytes)
    @mock.patch('ID2TLib.Utility.get_rnd_x86_nop', side_effect=Lib.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()