from unittest import mock
from routersploit.modules.exploits.routers.multi.tcp_32764_rce import Exploit


@mock.patch("routersploit.modules.exploits.routers.multi.tcp_32764_rce.shell")
def test_check_success1(mocked_shell, tcp_target):
    """ Test scenario - successful check Big Endian"""

    command_mock = tcp_target.get_command_mock(b"ABCDE")
    command_mock.return_value = b"MMcS"

    exploit = Exploit()
    exploit.target = tcp_target.host
    exploit.port = tcp_target.port

    assert exploit.check()
    assert exploit.run() is None


@mock.patch("routersploit.modules.exploits.routers.multi.tcp_32764_rce.shell")
def test_check_success2(mocked_shell, tcp_target):
    """ Test scenario - successful check - Little Endian"""

    command_mock = tcp_target.get_command_mock(b"ABCDE")
    command_mock.return_value = b"ScMM"

    exploit = Exploit()
    exploit.target = tcp_target.host
    exploit.port = tcp_target.port

    assert exploit.check()
    assert exploit.run() is None