diff --git a/conntrackt/views.py b/conntrackt/views.py --- a/conntrackt/views.py +++ b/conntrackt/views.py @@ -1,7 +1,6 @@ # For generating ZIP files. from StringIO import StringIO from zipfile import ZipFile, ZIP_DEFLATED -import re # Django-specific imports. from django.http import HttpResponse @@ -9,8 +8,10 @@ from django.shortcuts import render_to_r from django.views.generic import TemplateView, DetailView # Application-specific imports. -from conntrackt.models import Project, Entity, Location -from conntrackt.stripper import StripperMiddleware +from .models import Project, Entity, Location +from .stripper import StripperMiddleware +from .utils import generate_entity_iptables + class IndexView(TemplateView): """ @@ -35,6 +36,7 @@ class IndexView(TemplateView): return context + class ProjectView(DetailView): """ Custom view for presenting the project information. @@ -71,44 +73,64 @@ class ProjectView(DetailView): # Finally return the context. return context -class IptablesView(DetailView): + +class EntityView(DetailView): + """ + Custom view for presenting entity information. """ - Custom view for rendering iptables rules for a specific entity, and starting - automatic download by browser. + model = Entity + + def get_context_data(self, **kwargs): + """ + Returns the context data that should be used for rendering of the + template. + + Adds the 'entity_iptables' context object that contains full iptables + rules generated for the entity. + """ + + # Call the parent class method. + context = super(DetailView, self).get_context_data(**kwargs) + + # Add the rendered iptables rules to the context. + context['entity_iptables'] = generate_entity_iptables(self.object) + + return context + + +def entity_iptables(request, pk): + """ + Custom view that returns response containing iptables rules generated for an + entity. + + Makes sure to set the Content-Disposition of a response in order to + signal the browser it should start download of this view's response + immediately. Also sets the suggested filename for it. + + Arguments: + + pk - Primary key of the Entity object for which the rules should be + generated. + + Returns: + + Response object that contains the iptables rules for specified entity. """ - model = Entity - template_name = 'conntrackt/entity_iptables.html' - - def render_to_response(self, context, **kwargs): - """ - Renders the template with the given context to response. The response - will be an iptables configuration file that can be used with - iptables-restore. - - Makes sure to set the Content-Disposition of a response in order to - signal the browser it should start download of this view's response - immediately. Also sets the suggested filename for it. - """ + # Fetch the entity, and construct the response with iptables rules as + # content. + entity = get_object_or_404(Entity, pk = pk) + content = generate_entity_iptables(entity) + response = HttpResponse(content, mimetype='text/plain') + + # Add the Content-Disposition information for the browser, telling the + # browser to download the file with suggested filename. + response['Content-Disposition']="attachment; filename=%s-iptables.conf" % entity.name.lower().replace(" ", "_") - # Call the parent class method first. This will render the template. Set - # the content type to text/plain. - response = super(IptablesView, self).render_to_response(context, - content_type='text/plain', **kwargs) - - # Add the Content-Disposition information for the browser, telling the - # browser to download the file with suggested filename. - response['Content-Disposition']="attachment; filename=%s-iptables.conf" % self.object.name.lower().replace(" ", "_") + return response - # Render the response, and remove the blank lines from the template. - response.render() - response.content = re.sub('^\s*\n', '', response.content) - response.content = re.sub('\n\s*\n', '\n', response.content) - # Return the modified response. - return response - -def get_project_iptables(request, project_id, location_id = None): +def project_iptables(request, project_id, location_id = None): """ Custom view for obtaining iptables for all entities of a project or project location in a single ZIP file. @@ -143,10 +165,6 @@ def get_project_iptables(request, projec # to open the file with program as well. response = HttpResponse(mimetype='application/zip') - # Stripper middleware has to be used in order to remove the excess blank - # lines from rendering (this has to be done manually). - stripper_middleware = StripperMiddleware() - # If specific location was specified, get the entities that are part of that # project location only, otherwise fetch all of the project's entities. Also # set-up the filename that will be suggested to the browser. @@ -163,8 +181,8 @@ def get_project_iptables(request, projec # Render iptables rules for each entity, placing them in the ZIP archive. for entity in entities: - entity_iptables = stripper_middleware.process_response(request,render_to_response('conntrackt/entity_iptables.html', {'entity': entity}, mimetype="text/plain")) - zipped_iptables.writestr("%s-iptables.conf" % entity.name.lower().replace(" ", "_"), entity_iptables.content) + entity_iptables = generate_entity_iptables(entity) + zipped_iptables.writestr("%s-iptables.conf" % entity.name.lower().replace(" ", "_"), entity_iptables) # Close the archive, and flush the buffer. zipped_iptables.close()