source: src/django_gheat/gheat/management/commands/import_droidstumbler.py@ 9601

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Script for importing DroidStumbler .csv files, which takes the best value of
# each measurement point.
#
# Original by Dennis Wagenaar <d.wagenaar@gmail.com>
#
# Rick van der Zwet <info@rickvanderzwet.nl>
#
from django.core.management import setup_environ
from django.core.management.base import BaseCommand, CommandError
from django.db import connection, transaction
from django.db.utils import IntegrityError
from _mysql_exceptions import OperationalError
from optparse import OptionParser, make_option
from gheat.models import *
import csv
import datetime
import gzip
import logging
import os
import sys

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

def bulk_sql(sql_table, sql_values):
  if len(sql_values) == 0:
    raise ValueError, "No data to import"

  cursor = connection.cursor()
  try:
    # Make sure the special NULL is preserved
    sql = "INSERT INTO %s VALUES %s" % (sql_table, ','.join(sql_values).replace("'NULL'",'NULL'))
    count = cursor.execute(sql)
    transaction.commit_unless_managed()
  except OperationalError, e:
    logger.error("%s - %s ", sql_table, sql_values[0])
    raise
  except IntegrityError, e:
    logger.error("Unable to import - %s" %  e)
    raise
  return count

organizations = dict(Organization.objects.all().values_list('name','id'))
def get_organization_id_by_ssid(ssid):
  name = Organization.get_name_by_ssid(ssid)
  if not name:
    return 'NULL'
  else:
    return int(organizations[name])

def import_droidstumbler(filename, meetrondje):
  """ Import all points, return tuple with summary"""

  # Open file for reading
  if filename.endswith('.gz'):
    fh = gzip.open(filename,'rb')
  else:
    fh = open(filename,'rb')
  csvfile = csv.reader(fh, delimiter=',')

  #Various statistics
  counters = {'meting_added' : 0, 'meting_total' : 0, 'ap_added' : 0, 'ap_total' : 0}

  # Temponary holders
  meting_pool = {}
  ap_pool = {}
  # Process file, preparing new access points and measurements
  for row in csvfile:
    try:
      epoch, msg_type, lat, lon, accuracy, ssid, bssid, level, frequency, capabilities = row
      bssid = bssid.upper()
    except ValueError:
      # Known error, please ignore
      if row[1] == 'gps' and len(row) == 12: continue
      logger.error("Unable to parse line:%i '%s'" % (csvfile.line_num, row))
      continue
    if msg_type == "data" and lat and lon:
      counters['meting_total'] += 1
      if not ap_pool.has_key(bssid):
        encryption = 'WPA' in capabilities or 'WEP' in capabilities
        ap_pool[bssid] = (ssid, encryption)

      # We store the best value found
      key = (bssid, lat, lon)
      signaal=(100 + int(level))
      if meting_pool.has_key(key):
        meting_pool[key] = max(meting_pool[key], signaal)
      else:
        meting_pool[key] = signaal


  # Determine which entries we need to add
  counters['ap_total'] = len(ap_pool)
  bssid_list_present = Accespoint.objects.filter(mac__in=ap_pool.keys()).values_list('mac', flat=True)
  bssid_list_insert = set(ap_pool.keys()) - set(bssid_list_present)

  # Create a bulk import list and import
  if bssid_list_insert:
    sql_values = []
    for bssid in bssid_list_insert:
      ssid, encryption = ap_pool[bssid]
      # Special trick in SSID ts avoid escaping in later stage
      item = str((bssid.upper(),ssid.replace('%','%%'),encryption,get_organization_id_by_ssid(ssid)))
      sql_values.append(item)
    counters['ap_added'] = bulk_sql('gheat_accespoint (`mac`, `ssid`, `encryptie`, `organization_id`)',sql_values)

  # Build mapping for meting import
  mac2id = {}
  for mac,id in Accespoint.objects.filter(mac__in=ap_pool.keys()).values_list('mac','id'):
    mac2id[mac] = int(id)

  sql_values = []
  for (bssid,lat,lon),signal in meting_pool.iteritems():
    item = str((int(meetrondje.id),mac2id[bssid],float(lat),float(lon),int(signaal)))
    sql_values.append(item)

  # Bulk Import data if possible
  if sql_values:
    counters['meting_added'] = bulk_sql('gheat_meting (`meetrondje_id`, `accespoint_id`, `lat`, `lng`, `signaal`)',sql_values)
  return counters


class Command(BaseCommand):
  args = '<csvfile>[.gz] [csvfile2[.gz] [csvfile3[.gz] ...] '
  option_list = BaseCommand.option_list + (
    make_option('-m', '--meetrondje', dest='meetrondje', default=None),
    make_option('-g', '--gebruiker', dest='gebruiker', default=os.environ['USER']),
    make_option('-e', '--email', dest='email', default=os.environ['USER'] + '@example.org'),
    make_option('-k', '--kaart', dest='kaart', default='onbekend', help="Kaart gebruikt"),
    make_option('-d', '--datum', dest='datum', default=None, help="Provide date  \
      in following format: %Y-%m-%d-%H%M%S, by default it will be generated from \
      the filename"),
  )

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

    for csv_file in args:
      # Make sure to check files before we going to do importing at all
      if not os.path.isfile(csv_file):
        raise CommandError("csv file '%s' does not exists" % csv_file)
      logger.info("Processing '%s'" % csv_file)

      # Meetrondje from filename if needed
      if options['meetrondje'] == None:
        meetrondje = os.path.basename(csv_file).rstrip('.gz').rstrip('.csv') 
      else:
        meetrondje = options['meetrondje']
      # Date from filename if needed
      if options['datum'] == None:
         datum = os.path.basename(csv_file).lstrip('ScanResult-').rstrip('.csv.gz')
      else:
         datum = options['datum']
      try:
         datum = datetime.datetime.strptime(datum,'%Y-%m-%d-%H%M%S')
      except ValueError:
        raise CommandError("Invalid date '%s'" % options['datum'])

      # Create meetrondje object
      g, created = Gebruiker.objects.get_or_create(naam=options['gebruiker'] , email=options['email'])
      a, created = Apparatuur.objects.get_or_create(kaart=options['kaart'])
      mr, created = MeetRondje.objects.get_or_create(datum=datum , naam=meetrondje , gebruiker=g , apparatuur=a)
      logger.info('Meetrondje: %s @ %s' % (meetrondje, datum))
      if not created:
        logger.error("Meetrondje '%s' already imported" % mr)
        sys.exit(1)
      counters = import_droidstumbler(csv_file,mr)
      logger.info("summary accespoints: added:%(ap_added)-6s processed:%(ap_total)-6s" % counters)
      logger.info("summary metingen   : added:%(meting_added)-6s processed:%(meting_total)-6s" % counters)