Files @ e8434f637f9f
Branch filter:

Location: conntrackt/conntrackt/views.py

branko
Removed the line stripper middleware since it interfered with regular template lookups. Added some custom code to remove all blank lines from iptables rendering instead.
# For generating ZIP files.
from StringIO import StringIO
from zipfile import ZipFile, ZIP_DEFLATED
import re

# Django-specific imports.
from django.http import HttpResponse
from django.shortcuts import render_to_response, get_object_or_404
from django.views.generic import TemplateView, DetailView

# Application-specific imports.
from conntrackt.models import Project, Entity, Location
from conntrackt.stripper import StripperMiddleware

class IndexView(TemplateView):
    """
    Custom view used for rendering the index page.
    """

    template_name = 'conntrackt/index.html'

    def get_context_data(self, **kwargs):
        """
        Returns the context data that should be used for rendering of the
        template.

        Adds the 'projects' context object containing all of the projects
        available sorted aplhabetically by name.
        """

        # Set the context using the parent class.
        context = super(IndexView, self).get_context_data(**kwargs)

        context['projects'] = Project.objects.all().order_by('name')

        return context

class ProjectView(DetailView):
    """
    Custom view for presenting the project information.
    """

    model = Project

    def get_context_data(self, **kwargs):
        """
        Returns the context data that should be used for rendering of the
        template.

        Adds the 'location_entities' context object that contains tuples of the
        form '(location, entities)', where location is an instance of a
        Location, and entities is a query set of entities that belong to that
        particular location, and to related project.
        """

        # Set the context using the parent class.
        context = super(ProjectView, self).get_context_data(**kwargs)

        # Set-up an array that will contaion (location, entities) tuples.
        location_entities = []

        # Add the (location, entities) tuple for each location that has entities
        # belonging to the related project.
        for location in Location.objects.filter(entity__project = self.object).distinct():
            entities = Entity.objects.filter(project = self.object, location = location)
            location_entities.append((location, entities))

        # Add the (location, entities) tuples to context.
        context['location_entities'] = location_entities

        # Finally return the context.
        return context

class IptablesView(DetailView):
    """
    Custom view for rendering iptables rules for a specific entity, and starting
    automatic download by browser.
    """

    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.
        """

        # 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(" ", "_")

        # 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):
    """
    Custom view for obtaining iptables for all entities of a project or project
    location in a single ZIP file.

    Arguments:

        request - Request object.

        project_id - Unique ID of the project for whose entities the iptables
        rules should be generated.

        location_id - Optional unique ID of the project location for whose
        entities the iptables rules should be generated. Default is None, which
        means generate rules for _all_ entities in a project.

    Returns:

        Response object that contains the ZIP file and Content-Disposition
        information.
    """

    # Fetch the project.
    project = get_object_or_404(Project, pk = project_id)

    # Set-up a string IO object to which we'll write the ZIP file (in-memory).
    buff = StringIO()

    # Create a new ZIP file in-memory.
    zipped_iptables = ZipFile(buff, "w", ZIP_DEFLATED)

    # Create the response object, setting the mime type so browser could offer
    # 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.
    if location_id:
        location = get_object_or_404(Location, pk=location_id)
        entities = project.entity_set.filter(location = location)
        filename = '%s_%s-iptables.zip' % (project.name, location.name)
    else:
        entities = project.entity_set.all()
        filename = '%s-iptables.zip' % (project.name)

    # Lower-case the filename, and replace spaces with underscores (_).
    filename = filename.lower().replace(" ", "_")

    # 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)

    # Close the archive, and flush the buffer.
    zipped_iptables.close()
    buff.flush()

    # Write the contents of our buffer (ZIP archive) to response content, and
    # close the IO string.
    response.write(buff.getvalue())
    buff.close()

    # Set the Content-Disposition so the browser would know it should download
    # the archive, and suggest the filename.
    response['Content-Disposition'] = 'attachment; filename="%s"' % filename

    # Finally return the response object.
    return response