* Adding extra_headers feature

* Better handling for unicode
* Found pytries/datrie#20

moth/views/base/html_template_view.py

@@ -18,4 +18,10 @@ def get(self, request, *args, **kwds):
 
         context = self.get_context_data()
         context['html'] = self.HTML
-        return render(request, self.template_name, context)
+
+        response = render(request, self.template_name, context)
+
+        for key, value in self.extra_headers.iteritems():
+            response[key] = value
+
+        return response

moth/views/base/vulnerable_template_view.py

@@ -1,5 +1,6 @@
 import urlparse
 
+from django.shortcuts import render_to_response
 from django.views.generic import TemplateView
 from django.views.decorators.csrf import csrf_exempt
 from django.utils.decorators import method_decorator
@@ -40,11 +41,27 @@ class VulnerableTemplateView(TemplateView):
     linked = True
 
     plugin_families = set(get_plugin_families())
-
+
+    # Any extra HTTP response headers to add
+    extra_headers = {}
+
     @method_decorator(csrf_exempt)
     def dispatch(self, *args, **kwargs):
         return super(VulnerableTemplateView, self).dispatch(*args, **kwargs)
-
+
+    def get(self, request, *args, **kwds):
+        """
+        Adds extra headers to the response
+        """
+        # pylint: disable=E1101
+        context = self.get_context_data()
+
+        response = render_to_response(self.template_name, context)
+        for key, value in self.extra_headers.iteritems():
+            response[key] = value
+
+        return response
+
     def get_context_data(self, **kwargs):
         context = super(VulnerableTemplateView, self).get_context_data(**kwargs)
         context['title'] = self.title
@@ -63,7 +80,9 @@ def get_url_path(self):
         family, plugin = self.get_family_plugin()
         path = urlparse.urlparse(self.url_path).path
 
-        return unicode('%s/%s/%s' % (family, plugin, path))
+        url_path = '%s/%s/%s' % (family, plugin, path)
+
+        return url_path.decode('utf-8')
 
     def get_family_plugin(self):
         '''

moth/views/router.py

@@ -28,7 +28,11 @@ class RouterView(object):
 
     def __init__(self):
         self._plugin_families = set(get_plugin_families())
-        self._mapping = Trie(string.printable)
+
+        # All chars are allowed, this is required for links with unicode chars
+        allowed_charset = [chr(x) for x in range(0, 255)]
+        self._mapping = Trie(allowed_charset)
+
         self._view_files = []
         self._autoregister()
 
@@ -192,7 +196,7 @@ def __call__(self, request, *args, **kwargs):
         else:
             # Try to create an "Index of" page for vulnerabilities
             sub_views = self._mapping.values(url_path)
-            
+
             if sub_views:
                 return self._generate_index(request, url_path, sub_views)