source: src/django_gheat/gheat/management/commands/import_datafile.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Script for importing various stumble files.
#
# Rick van der Zwet <info@rickvanderzwet.nl>
#
from django.core.files import File
from django.core.management.base import BaseCommand,CommandError
from django.db import connection, transaction
from django.db.utils import IntegrityError
from gheat.models import *
from optparse import OptionParser, make_option
import datetime
import gzip
import os
import sys
import logging

from collections import defaultdict

from gheat.dataimport import import_file

logging.basicConfig(level=logging.INFO, format='# %(levelname)s: %(message)s')
logger = logging.getLogger()


valid_prefixes = ['DroidStumbler-', 'Kismet-EeePC-', 'Kismet-', 'ScanResult-']
def strip_prefix(filename):
  """ Prefix removal """
  for prefix in valid_prefixes:
    if filename.startswith(prefix):
      filename = filename[len(prefix):]
  return filename


valid_suffixes = ['.gz', '.gpsxml', '.netxml', '.csv', '.ns1']
def strip_suffix(filename):
  """ Suffix removal """
  for suffix in valid_suffixes:
    if filename.endswith(suffix):
      filename = filename[:-len(suffix)]
  return filename


def strip_file(filename):
  """ Prefix and suffix removal """
  return strip_suffix(strip_prefix(filename))


#Kismet-20110805-15-37-30-1
#ScanResult-2011-05-09-201117
strptime_choices = ['%Y%m%d-%H-%M-%S-1', '%Y-%m-%d-%H%M%S']
def process_date(datestr):
  for strptime in strptime_choices:
    try:
      return datetime.datetime.strptime(datestr,strptime)
    except ValueError:
      pass
  logger.error("Invalid date '%s', options: %s, using: now()", datestr, strptime_choices)
  return datetime.datetime.now()


class Command(BaseCommand):
  args = '<netstumber.ns1>[.gz] [netstumber2.ns1[.gz]  netstumber3.ns1[.gz] ...]'
  option_list = BaseCommand.option_list + (
    make_option('-k', '--kaart', dest='kaart', default='onbekend', 
      help="Kaart gebruikt"),
    make_option('-f', '--force', dest='force', default=False, action="store_true",
      help="Import anyways, even if the meetrondje is already imported"),
    make_option('-m', '--meetrondje', dest='meetrondje', default=None),
    make_option('-g', '--gebruiker', dest='gebruiker', default='username',
      help='Naam van de persoon die de meting uitgevoerd heeft'),
    make_option('-e', '--email', dest='email', default='foo@bar.org',
      help='Email van de persoon die de meting uitgevoerd heeft'),
    make_option('-d', '--datum', dest='datum', default=None,
      help="Provide date in following format: '%Y%m%d-%H-%M-%S-1', by \
      default it will be generated from the filename"),
    make_option('-b', '--batch', dest='batch', default=False, action="store_true",
      help="Batch import modes, gives no errors if items is already imported"),
  )

  def handle(self, *args, **options):
    if options['verbosity'] == 2:
      logger.setLevel(logging.DEBUG)
    if len(args) == 0:
      self.print_help(sys.argv[0],sys.argv[1])
      raise CommandError("Not all arguments are provided")

    # Please first the netxml and the gpsxml files and the rest
    sorted_args = [x for x in args if "netxml" in x] +\
     [x for x in args if "gpsxml" in x] +\
     [x for x in args if "ns1" in x]
    remainder = list(set(args) - set(sorted_args))
    args = sorted_args + remainder
    logger.debug("Parsing files in the following order: %s", args)

    # Make sure the all exists at first
    for filename in args:
      if not os.path.isfile(filename):
        raise CommandError("file '%s' does not exists" % filename)


    def get_date(filename):
      if options['datum'] == None:
         datestr = strip_file(os.path.basename(filename))
         datum = process_date(datestr)
      elif options['datum'] == 'now':
         datum = datetime.datetime.now()
      else:
         datum = process_date(options['datum'])
      return datum

    def get_meetrondje(meetrondje):
      # Meetrondje from filename if needed
      if options['meetrondje'] == None:
        meetrondje = strip_suffix(os.path.basename(filename))
      else:
        meetrondje = options['meetrondje']
      return meetrondje

    # Get Gheat Objects, pre-req
    gebruiker, created = Gebruiker.objects.get_or_create(naam=options['gebruiker'],
      email=options['email'])
    apparatuur, created = Apparatuur.objects.get_or_create(kaart=options['kaart'])

    # Meetrondje is deducted and checked from first filename
    filename = args[0]
    logger.info("Processing '%s'" % filename)
    meetrondje, created = MeetRondje.objects.get_or_create(
      datum=get_date(filename), naam=get_meetrondje(filename),
      gebruiker=gebruiker, apparatuur=apparatuur)
    if not options['force'] and not created:
      raise CommandError("Meetrondje '%s' already imported",  meetrondje)

    # Check if all files are valid
    for filename in args:
      logger.info("Meetrondje: %s", meetrondje)
      logger.info("Bestand: %s", filename)
      meetbestand = MeetBestand(meetrondje=meetrondje,is_imported=True)
      meetbestand.bestand.save(os.path.basename(filename),File(open(filename)))
      meetbestand.save()

      counters = import_file(filename,meetrondje)
      logger.info("summary accesspoints: total:%(ap_total)-6s added:%(ap_added)-6s failed:%(ap_failed)-6s ignored:%(ap_ignored)-6s" % counters)
      logger.info("summary client     : total:%(client_total)-6s added:%(client_added)-6s failed:%(client_failed)-6s ignored:%(client_ignored)-6s" % counters)
      logger.info("summary metingen   : total:%(meting_total)-6s added:%(meting_added)-6s failed:%(meting_failed)-6s ignored:%(meting_ignored)-6s" % counters)