from drozer.modules import Module

class Vulnerability(object):
    """
    A special module that contains a test and exploit for a particular vulnerability.

    Modules that use this library do not need to define an execute() method,
    and may be combined with a VulnerabilityScanner.

    NOT_VULNERABLE = 0
    VULNERABLE = 1
    POSSIBLY_VULNERABLE = 2
    """

    _vulnerability_definition = True

    def execute(self, arguments):
        """
        Implementation of the Module execute() method. If the device is vulnerable,
        the exploit is run. Otherwise, 'Not Vulnerable' is printed to screen.
        """

        if self.isVulnerable(arguments):
            self.exploit(arguments)
        else:
            self.stdout.write("Not Vulnerable.\n")

class VulnerabilityScanner(object):
    """
    A special scanner module, that finds all vulnerability definitions under the
    path specified in 'vulnerabilities' and runs the checks in turn.
    """

    def execute(self, arguments):
        """
        Implementation of the Module execute() method. Each included vulnerability
        test is executed in turn, and printed out as 'Vulnerable', 'Not Vulnerable',
        or 'Exception'.
        """

        for vulnerability in self.__vulnerabilities():
            try:
                check = self.modules.get(vulnerability)(self)
                label = hasattr(check, 'label') and check.label or check.fqmn()
                result = check.isVulnerable(arguments)

                self.stdout.write("%s - %s\n" % (label, result==1 and "[color red]Vulnerable[/color]" or (result==2 and "[color yellow]Possibly Vulnerable[/color]") or "[color green]Not Vulnerable[/color]"))
            except:
                self.stdout.write("%s - %s\n" % (label, "[color blue]Exception[/color]"))

    def __vulnerabilities(self):
        """
        Find all vulnerability definition modules under 'vulnerabilities'.
        """
        
        return filter(lambda m: m.startswith(self.vulnerabilities) and hasattr(self.modules.get(m), '_vulnerability_definition') and self.modules.get(m)._vulnerability_definition, self.modules.all())
        
