def test_iterkeys(self):
        disk_dict = DiskDict()

        disk_dict['a'] = 'abc'
        disk_dict['b'] = 'abc'
        disk_dict['c'] = 'abc'

        self.assertEqual(set(disk_dict.iterkeys()), set(['a', 'b', 'c']))

    def __init__(self):

        self._variants_eq = DiskDict(table_prefix='variant_db_eq')
        self._variants = DiskDict(table_prefix='variant_db')

        self.params_max_variants = cf.cf.get('params_max_variants', PARAMS_MAX_VARIANTS)
        self.path_max_variants = cf.cf.get('path_max_variants', PATH_MAX_VARIANTS)

        self._db_lock = threading.RLock()

    def __init__(self):
        self._variants = DiskDict(table_prefix='variant_db')
        self._variants_eq = DiskDict(table_prefix='variant_db_eq')
        self._variants_form = DiskDict(table_prefix='variant_db_form')

        self.params_max_variants = cf.cf.get('params_max_variants')
        self.path_max_variants = cf.cf.get('path_max_variants')
        self.max_equal_form_variants = cf.cf.get('max_equal_form_variants')

        self._db_lock = threading.RLock()

 def test_remove_table(self):
     disk_dict = DiskDict()
     table_name = disk_dict.table_name
     db = get_default_temp_db_instance()
     
     self.assertTrue(db.table_exists(table_name))
     
     disk_dict.cleanup()
     
     self.assertFalse(db.table_exists(table_name))

    def test_table_with_prefix(self):
        _unittest = 'unittest'
        disk_dict = DiskDict(_unittest)

        self.assertIn(_unittest, disk_dict.table_name)
        db = get_default_temp_db_instance()

        self.assertTrue(db.table_exists(disk_dict.table_name))

        disk_dict.cleanup()

        self.assertFalse(db.table_exists(disk_dict.table_name))

    def test_get(self):
        disk_dict = DiskDict()

        disk_dict[0] = 'abc'
        
        abc1 = disk_dict.get(0)
        abc2 = disk_dict.get(0, 1)
        two = disk_dict.get(1, 2)
        
        self.assertEqual(abc1, 'abc')
        self.assertEqual(abc2, 'abc')
        self.assertEqual(two, 2)

    def __init__(self):
        GrepPlugin.__init__(self)

        # Internal variables
        self._comments = DiskDict(table_prefix='html_comments')
        self._already_reported = ScalableBloomFilter()
        self._end_was_called = False

    def __init__(self):
        AuditPlugin.__init__(self)

        # Internal variables
        self._persistent_multi_in = None
        self._expected_mutant_dict = DiskDict(table_prefix='ssi')
        self._extract_expected_re = re.compile('[1-9]{5}')

 def _init(self, maxsize):
     """
     Initialize the dicts and pointer
     :param maxsize: The max size for the queue
     """
     self.queue_order = list()
     self.hash_to_uuid = dict()
     self.memory = dict()
     self.disk = DiskDict(table_prefix='%sCachedQueue' % self.name)

 def _init(self, maxsize):
     """
     Initialize the dicts and pointer
     :param maxsize: The max size for the queue
     """
     self.memory = dict()
     self.disk = DiskDict(table_prefix='%sCachedQueue' % self.name)
     self.get_pointer = 0
     self.put_pointer = 0

    def __init__(self):
        AuditPlugin.__init__(self)

        # Internal variables
        self._expected_res_mutant = DiskDict()
        self._freq_list = DiskList()
        
        re_str = '<!--#exec cmd="echo -n (.*?);echo -n (.*?)" -->'
        self._extract_results_re = re.compile(re_str) 

    def __init__(self, max_in_memory=50, table_prefix=None):
        """
        :param max_in_memory: The max number of items to keep in memory
        """
        assert max_in_memory > 0, 'In-memory items must be > 0'

        table_prefix = self._get_table_prefix(table_prefix)

        self._max_in_memory = max_in_memory
        self._disk_dict = DiskDict(table_prefix=table_prefix)
        self._in_memory = dict()
        self._access_count = dict()

class ssi(AuditPlugin):
    """
    Find server side inclusion vulnerabilities.
    :author: Andres Riancho ([email protected])
    """

    def __init__(self):
        AuditPlugin.__init__(self)

        # Internal variables
        self._expected_res_mutant = DiskDict()
        self._freq_list = DiskList()
        
        re_str = '<!--#exec cmd="echo -n (.*?);echo -n (.*?)" -->'
        self._extract_results_re = re.compile(re_str) 

    def audit(self, freq, orig_response):
        """
        Tests an URL for server side inclusion vulnerabilities.

        :param freq: A FuzzableRequest
        """
        # Create the mutants to send right now,
        ssi_strings = self._get_ssi_strings()
        mutants = create_mutants(freq, ssi_strings, orig_resp=orig_response)

        # Used in end() to detect "persistent SSI"
        for mut in mutants:
            expected_result = self._extract_result_from_payload(
                mut.get_token_value())
            self._expected_res_mutant[expected_result] = mut

        self._freq_list.append(freq)
        # End of persistent SSI setup

        self._send_mutants_in_threads(self._uri_opener.send_mutant,
                                      mutants,
                                      self._analyze_result)

    def _get_ssi_strings(self):
        """
        This method returns a list of server sides to try to include.

        :return: A string, see above.
        """
        yield '<!--#exec cmd="echo -n %s;echo -n %s" -->' % (rand_alpha(5),
                                                             rand_alpha(5))

        # TODO: Add mod_perl ssi injection support
        # http://www.sens.buffalo.edu/services/webhosting/advanced/perlssi.shtml
        #yield <!--#perl sub="sub {print qq/If you see this, mod_perl is working!/;}" -->

    def _extract_result_from_payload(self, payload):
        """
        Extract the expected result from the payload we're sending.
        """
        match = self._extract_results_re.search(payload)
        return match.group(1) + match.group(2)

    def _analyze_result(self, mutant, response):
        """
        Analyze the result of the previously sent request.
        :return: None, save the vuln to the kb.
        """
        if self._has_no_bug(mutant):
            e_res = self._extract_result_from_payload(mutant.get_token_value())
            if e_res in response and not e_res in mutant.get_original_response_body():
                
                desc = 'Server side include (SSI) was found at: %s'
                desc = desc % mutant.found_at()
                
                v = Vuln.from_mutant('Server side include vulnerability', desc,
                                     severity.HIGH, response.id, self.get_name(),
                                     mutant)

                v.add_to_highlight(e_res)
                self.kb_append_uniq(self, 'ssi', v)

    def end(self):
        """
        This method is called when the plugin wont be used anymore and is used
        to find persistent SSI vulnerabilities.

        Example where a persistent SSI can be found:

        Say you have a "guestbook" (a CGI application that allows visitors
        to leave messages for everyone to see) on a server that has SSI
        enabled. Most such guestbooks around the Net actually allow visitors
        to enter HTML code as part of their comments. Now, what happens if a
        malicious visitor decides to do some damage by entering the following:

        <!--#exec cmd="ls" -->

        If the guestbook CGI program was designed carefully, to strip SSI
        commands from the input, then there is no problem. But, if it was not,
        there exists the potential for a major headache!

        For a working example please see moth VM.
        """
        multi_in_inst = multi_in(self._expected_res_mutant.keys())
#.........这里部分代码省略.........

class CachedQueue(Queue.Queue, QueueSpeedMeasurement):
    """
    The framework uses the producer / consumer design pattern extensively.
    In order to avoid high memory usage in the queues connecting the different
    parts of the framework we defined a max size.

    When a queue max size is reached, one or more threads will block. This
    line is printed during a real scan:

        Thread blocked 5.76617312431 seconds waiting for Queue.put() to have space
        in the Grep queue. The queue's maxsize is 20.

    In the case of the Grep consumer / producer the problem with a block is increased
    by the fact that HTTP responses won't reach other parts of the framework
    until the queue has space.

    Increasing the queue size would increase memory usage.

    Using an on-disk queue would increase CPU (serialization) and disk IO.

    The CacheQueue is a mix of in-memory and on-disk queue. The first N items
    are stored in memory, when more items are put() we just write them to
    disk.

    The CacheQueue object implements these methods from QueueSpeedMeasurement:
        * get_input_rpm
        * get_output_rpm

    Which allows users to understand how fast a queue is moving.
    """
    def __init__(self, maxsize=0, name='Unknown'):
        self.name = name
        self.max_in_memory = maxsize

        QueueSpeedMeasurement.__init__(self)

        # We want to send zero to the maxsize of the Queue implementation
        # here because we can write an infinite number of items
        Queue.Queue.__init__(self, maxsize=0)

    def get_name(self):
        return self.name

    def next_item_saved_to_memory(self):
        return len(self.memory) < self.max_in_memory

    def _init(self, maxsize):
        """
        Initialize the dicts and pointer
        :param maxsize: The max size for the queue
        """
        self.memory = dict()
        self.disk = DiskDict(table_prefix='%sCachedQueue' % self.name)
        self.get_pointer = 0
        self.put_pointer = 0

    def _qsize(self, len=len):
        return len(self.memory) + len(self.disk)

    def _get_class_name(self, obj):
        try:
            return obj.__class__.__name__
        except:
            return type(obj)

    def _put(self, item):
        """
        Put a new item in the queue
        """
        #
        #   This is very useful information for finding bottlenecks in the
        #   framework / strategy
        #
        if len(self.memory) == self.max_in_memory:
            #
            #   If you see many messages like this in the scan log, then you
            #   might want to experiment with a larger maxsize for this queue
            #
            msg = ('CachedQueue.put() will write a %r item to the %s DiskDict.'
                   ' This uses more CPU and disk IO than storing in memory'
                   ' but will avoid high memory usage issues. The current'
                   ' %s DiskDict size is %s.')
            args = (self._get_class_name(item),
                    self.get_name(),
                    self.get_name(),
                    len(self.disk))
            om.out.debug(msg % args)

        #
        #   And now we just save the item to memory (if there is space) or
        #   disk (if it doesn't fit on memory)
        #
        if len(self.memory) < self.max_in_memory:
            self.memory[self.put_pointer] = item
        else:
            self.disk[self.put_pointer] = item

        self.put_pointer += 1
        self._item_added_to_queue()

#.........这里部分代码省略.........

class html_comments(GrepPlugin):
    """
    Extract and analyze HTML comments.

    :author: Andres Riancho ([email protected])
    """

    HTML_RE = re.compile('<[a-zA-Z]*.*?>.*?</[a-zA-Z]>')

    INTERESTING_WORDS = (
        # In English
        'user', 'pass', 'xxx', 'fix', 'bug', 'broken', 'oops', 'hack',
        'caution', 'todo', 'note', 'warning', '!!!', '???', 'shit',
        'pass', 'password', 'passwd', 'pwd', 'secret', 'stupid',
        
        # In Spanish
        'tonto', 'porqueria', 'cuidado', 'usuario', u'contraseña',
        'puta', 'email', 'security', 'captcha', 'pinga', 'cojones',
        
        # some in Portuguese
        'banco', 'bradesco', 'itau', 'visa', 'bancoreal', u'transfêrencia',
        u'depósito', u'cartão', u'crédito', 'dados pessoais'
    )

    _multi_in = MultiIn([' %s ' % w for w in INTERESTING_WORDS])

    def __init__(self):
        GrepPlugin.__init__(self)

        # Internal variables
        self._comments = DiskDict(table_prefix='html_comments')
        self._already_reported = ScalableBloomFilter()
        self._end_was_called = False

    def grep(self, request, response):
        """
        Plugin entry point, parse those comments!

        :param request: The HTTP request object.
        :param response: The HTTP response object
        :return: None
        """
        if not response.is_text_or_html():
            return
        
        try:
            dp = parser_cache.dpc.get_document_parser_for(response)
        except BaseFrameworkException:
            return
        
        for comment in dp.get_comments():
            # These next two lines fix this issue:
            # audit.ssi + grep.html_comments + web app with XSS = false positive
            if request.sent(comment):
                continue

            if self._is_new(comment, response):

                self._interesting_word(comment, request, response)
                self._html_in_comment(comment, request, response)

    def _interesting_word(self, comment, request, response):
        """
        Find interesting words in HTML comments
        """
        comment = comment.lower()

        for word in self._multi_in.query(comment):
            if (word, response.get_url()) in self._already_reported:
                continue

            desc = ('A comment with the string "%s" was found in: "%s".'
                    ' This could be interesting.')
            desc %= (word, response.get_url())

            i = Info.from_fr('Interesting HTML comment', desc, response.id,
                             self.get_name(), request)
            i.add_to_highlight(word)

            kb.kb.append(self, 'interesting_comments', i)
            om.out.information(i.get_desc())
                
            self._already_reported.add((word, response.get_url()))

    def _html_in_comment(self, comment, request, response):
        """
        Find HTML code in HTML comments
        """
        html_in_comment = self.HTML_RE.search(comment)

        if html_in_comment is None:
            return

        if (comment, response.get_url()) in self._already_reported:
            return

        # There is HTML code in the comment.
        comment = comment.strip()
        comment = comment.replace('\n', '')
        comment = comment.replace('\r', '')
#.........这里部分代码省略.........

class ssi(AuditPlugin):
    """
    Find server side inclusion vulnerabilities.
    :author: Andres Riancho ([email protected])
    """

    def __init__(self):
        AuditPlugin.__init__(self)

        # Internal variables
        self._expected_mutant_dict = DiskDict(table_prefix="ssi")
        self._extract_expected_re = re.compile("[1-9]{5}")

    def audit(self, freq, orig_response):
        """
        Tests an URL for server side inclusion vulnerabilities.

        :param freq: A FuzzableRequest
        """
        ssi_strings = self._get_ssi_strings()
        mutants = create_mutants(freq, ssi_strings, orig_resp=orig_response)

        self._send_mutants_in_threads(self._uri_opener.send_mutant, mutants, self._analyze_result)

    def _get_ssi_strings(self):
        """
        This method returns a list of server sides to try to include.

        :return: A string, see above.
        """
        # Generic
        yield '<!--#exec cmd="echo -n %s;echo -n %s" -->' % get_seeds()

        # Perl SSI
        yield (
            '<!--#set var="SEED_A" value="%s" -->'
            '<!--#echo var="SEED_A" -->'
            '<!--#set var="SEED_B" value="%s" -->'
            '<!--#echo var="SEED_B" -->' % get_seeds()
        )

        # Smarty
        # http://www.smarty.net/docsv2/en/language.function.math.tpl
        yield '{math equation="x * y" x=%s y=%s}' % get_seeds()

        # Mako
        # http://docs.makotemplates.org/en/latest/syntax.html
        yield "${%s * %s}" % get_seeds()

        # Jinja2 and Twig
        # http://jinja.pocoo.org/docs/dev/templates/#math
        # http://twig.sensiolabs.org/doc/templates.html
        yield "{{%s * %s}}" % get_seeds()

        # Generic
        yield "{%s * %s}" % get_seeds()

    def _get_expected_results(self, mutant):
        """
        Extracts the potential results from the mutant payload and returns them
        in a list.
        """
        sent_payload = mutant.get_token_payload()
        seed_numbers = self._extract_expected_re.findall(sent_payload)

        seed_a = int(seed_numbers[0])
        seed_b = int(seed_numbers[1])

        return [str(seed_a * seed_b), "%s%s" % (seed_a, seed_b)]

    def _analyze_result(self, mutant, response):
        """
        Analyze the result of the previously sent request.
        :return: None, save the vuln to the kb.
        """
        # Store the mutants in order to be able to analyze the persistent case
        # later
        expected_results = self._get_expected_results(mutant)

        for expected_result in expected_results:
            self._expected_mutant_dict[expected_result] = mutant

        # Now we analyze the "reflected" case
        if self._has_bug(mutant):
            return

        for expected_result in expected_results:
            if expected_result not in response:
                continue

            if expected_result in mutant.get_original_response_body():
                continue

            desc = "Server side include (SSI) was found at: %s"
            desc %= mutant.found_at()

            v = Vuln.from_mutant(
                "Server side include vulnerability", desc, severity.HIGH, response.id, self.get_name(), mutant
            )

#.........这里部分代码省略.........

 def __init__(self, max_variants=5):
     self._disk_dict = DiskDict()
     self._db_lock = threading.RLock()
     self.max_variants = max_variants

    def __init__(self):
        AuditPlugin.__init__(self)

        # Internal variables
        self._expected_mutant_dict = DiskDict(table_prefix="ssi")
        self._extract_expected_re = re.compile("[1-9]{5}")

 def __init__(self, iterable=(), maxsize=-1):
     if not hasattr(self, 'data'):
         self.left = self.right = 0
         self.data = DiskDict()
     self.maxsize = maxsize
     self.extend(iterable)

class VariantDB(object):

    def __init__(self, max_variants=DEFAULT_MAX_VARIANTS):
        self._disk_dict = DiskDict(table_prefix='variant_db')
        self._db_lock = threading.RLock()
        self.max_variants = max_variants

    def append(self, reference):
        """
        Called when a new reference is found and we proved that new
        variants are still needed.

        :param reference: The reference (as a URL object) to add. This method
                          will "normalize" it before adding it to the internal
                          shelve.
        """
        clean_reference = self._clean_reference(reference)

        with self._db_lock:
            count = self._disk_dict.get(clean_reference, None)

            if count is not None:
                self._disk_dict[clean_reference] = count + 1
            else:
                self._disk_dict[clean_reference] = 1

    def need_more_variants(self, reference):
        """
        :return: True if there are not enough variants associated with
        this reference in the DB.
        """
        clean_reference = self._clean_reference(reference)

        # I believe this is atomic enough...
        count = self._disk_dict.get(clean_reference, 0)
        if count >= self.max_variants:
            return False
        else:
            return True

    def _clean_reference(self, reference):
        """
        This method is VERY dependent on the are_variants method from
        core.data.request.variant_identification , make sure to remember that
        when changing stuff here or there.

        What this method does is to "normalize" any input reference string so
        that they can be compared very simply using string match.

        """
        res = reference.get_domain_path() + reference.get_file_name()

        if reference.has_query_string():

            res += '?'
            qs = copy.deepcopy(reference.querystring)

            for key, value, path, setter in qs.iter_setters():

                if value.isdigit():
                    setter('number')
                else:
                    setter('string')

            res += str(qs)

        return res