from io import StringIO
import sys
from unittest.mock import patch, MagicMock
import malwarebazaar
import unittest

class TestMalwarebazaarMethods(unittest.TestCase):
    def test_main_missing_input(self):
        with patch('sys.stdout', new=StringIO()) as mock_cmd:
            sys.argv = ["cmd"]
            malwarebazaar.main()
            self.assertEqual(mock_cmd.getvalue(),
                             'ERROR: Input is not in proper JSON format\n')

    def test_main_success(self):
            with patch('sys.stdout', new=StringIO()) as mock_cmd:
                with patch('malwarebazaar.analyze', new=MagicMock(return_value={'test': 'val'})) as mock:
                    sys.argv = ["cmd", "input"]
                    malwarebazaar.main()
                    expected = '{"test": "val"}\n'
                    self.assertEqual(mock_cmd.getvalue(), expected)
                    mock.assert_called_once()
    def test_analyze(self):
        """simulated sendReq and prepareResults with 2 mock objects and variables sendReqOutput and prepareResultOutput,
            input created for analyze method call and then we compared results['summary'] with 'no result' """
        sendReqOutput = {'threat': 'no_result',"query_status":"ok",'data':[{'sha256_hash':'notavalidhash'}]}
        input = '{"artifactType":"hash", "value":"1234"}'
        input2 ='{"artifactType":"tlsh", "value":"1234"}'
        input3='{"artifactType":"gimphash", "value":"1234"}'
        prepareResultOutput = {'response': '',
                               'summary': 'no result', 'status': 'info'}

        with patch('malwarebazaar.sendReq', new=MagicMock(return_value=sendReqOutput)) as mock:
            with patch('malwarebazaar.prepareResults', new=MagicMock(return_value=prepareResultOutput)) as mock2:
                results = malwarebazaar.analyze(input)
                results2 = malwarebazaar.analyze(input2)
                results3 =malwarebazaar.analyze(input3)
                self.assertEqual(results["summary"],prepareResultOutput['summary'])
                self.assertEqual(results2["summary"], prepareResultOutput['summary'])
                self.assertEqual(results3["summary"], prepareResultOutput['summary'])
                self.assertEqual(results["status"], "info")
                self.assertEqual(results2["status"], "info")
                self.assertEqual(results3["status"], "info")

                mock.assert_called()

    def test_prepareResults_illegal_search_term(self):
            # illegal search term
            raw = {'query_status': 'illegal_search_term'}
            expected = {'response': raw, 'status': 'info', 'summary': 'no result'}
            results = malwarebazaar.prepareResults(raw)
            self.assertEqual(results, expected)

    def test_buildReqGimqhash(self):
        result = malwarebazaar.buildReq('gimphash', '')
        self.assertEqual(
            result, {'query': 'get_gimphash', 'gimphash': ''})
        
    def test_buildReqHash(self):
        result = malwarebazaar.buildReq('hash', '')
        self.assertEqual(
            result, {'query': 'get_info', 'hash': ''})
    def test_buildReqtlshhash(self):
        result = malwarebazaar.buildReq('tlsh', '')
        self.assertEqual(
            result, {'query': 'get_tlsh', 'tlsh': ''})