app.py - vendor/google/wavedroplet - Git at Google

 # Copyright 2014 Google Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 # http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.

 """Handlers for uploading, filtering, and visualization."""

 import collections
 import errno
 import json
 import sys
 import time
 import traceback
 import urllib
 import webapp2
 import tornado.template
 import wifipacket
 from google.appengine.api import memcache
 from google.appengine.api import users
 from google.appengine.ext import blobstore
 from google.appengine.ext import ndb
 from google.appengine.ext.webapp import blobstore_handlers

 BROADCAST = 'ff:ff:ff:ff:ff:ff'
 DEFAULT_FIELDS = ['seq', 'rate']
 AVAIL_FIELDS = ['seq', 'mcs', 'spatialstreams', 'bw', 'rate', 'retry',
                 'type', 'typestr', 'dbm_antsignal', 'dbm_antnoise',
                 'bad']

 IS_DEBUG = False
 SAMPLE_SIZE = 2

 loader = tornado.template.Loader('.')


 def _Esc(s):
   """Like tornado.escape.url_escape, but only escapes &, #, %, and =."""
   out = []
   for c in s:
     if c in ['&', '#', '%', '=']:
       out.append('%%%02X' % ord(c))
     else:
       out.append(c)
   return ''.join(out)


 def AllowedEmails():
   try:
     return open('email-allow.txt').read().split()
   except IOError as e:
     if e.errno == errno.ENOENT:
       pass
     else:
       raise
   return []


 def GoogleLoginRequired(func):
   """Enforcing @google.com login."""

   def Handler(self, *args, **kwargs):
     user = users.get_current_user()
     if not user:
       self.redirect(users.create_login_url('/'))
     elif (not user.email().endswith('@google.com') and
           user.email() not in AllowedEmails()):
       self.response.set_status(401, 'Unauthorized')
       self.response.write("Sorry.  You're not an authorized user.")
     else:
       return func(self, *args, **kwargs)
   return Handler


 class PcapData(ndb.Model):
   """Info about pcap file and its visualization settings/cache."""

   create_time = ndb.DateTimeProperty(auto_now_add=True)
   create_user_email = ndb.StringProperty()
   filename = ndb.StringProperty()
   show_hosts = ndb.StringProperty(repeated=True)
   show_fields = ndb.StringProperty(repeated=True)
   aliases = ndb.PickleProperty()

   # Cached JSON representations for various useful subsets of data
   # to be passed to the JS side for visualization.
   # All captured packets
   js_packets = ndb.JsonProperty(compressed=True)
   # All pairs of (transmitter, receiver)
   js_streams = ndb.JsonProperty(compressed=True)

   @staticmethod
   def GetDefault():
     return PcapData.get_or_insert(str('*'), show_hosts=[], aliases={})

   @staticmethod
   def GetOrInsertFromBlob(blob_info):
     u = users.get_current_user()
     if u:
       email = u.email()
     else:
       email = '<anonymous>'
     return PcapData.get_or_insert(str(blob_info.key()),
                                   show_hosts=[], aliases={},
                                   filename=blob_info.filename,
                                   create_user_email=email)


 class _BaseHandler(webapp2.RequestHandler):

   def render(self, template, **kwargs):
     d = dict()
     d.update(kwargs)
     self.response.write(loader.load(template).generate(**d))


 class MainHandler(_BaseHandler):

   @GoogleLoginRequired
   def get(self):
     upload_url = blobstore.create_upload_url('/upload')
     q = PcapData.query().order(-PcapData.create_time).fetch(10)
     self.render('index.html',
                 upload_url=upload_url,
                 recents=q)


 class UploadHandler(blobstore_handlers.BlobstoreUploadHandler):

   def post(self):
     upload_files = self.get_uploads()
     sys.stderr.write('upload: %r\n' % upload_files)
     blob_info = upload_files[0]
     reader = blob_info.open()
     try:
       wifipacket.Packetize(reader).next()  # just check basic file header
     except wifipacket.Error:
       blob_info.delete()
       raise
     self.redirect('/view/%s' % blob_info.key())


 class DownloadHandler(blobstore_handlers.BlobstoreDownloadHandler):

   def get(self, blobres):
     blob_info = blobstore.BlobInfo.get(str(urllib.unquote(blobres)))
     self.send_blob(blob_info)


 def _Boxes(blob_info):
   """Re-/store from/to memcache number of packets per mac address."""

   boxes = memcache.get(str(blob_info.key()), namespace='boxes')
   if not boxes:
     reader = blob_info.open()
     boxes = collections.defaultdict(lambda: 0)
     # TODO(katepek): use cache here instead if available
     for p, unused_frame in wifipacket.Packetize(reader):
       if 'flags' in p and p.flags & wifipacket.Flags.BAD_FCS: continue
       if 'ta' in p and 'ra' in p:
         boxes[p.ta] += 1
         boxes[p.ra] += 1
     memcache.add(key=str(blob_info.key()), value=dict(boxes),
                  namespace='boxes')
   return boxes


 class ViewHandler(_BaseHandler):

   @GoogleLoginRequired
   def get(self, blobres):
     blob_info = blobstore.BlobInfo.get(str(urllib.unquote(blobres)))
     capdefault = PcapData.GetDefault()
     pcapdata = PcapData.GetOrInsertFromBlob(blob_info)

     boxes = _Boxes(blob_info)
     cutoff = max(boxes.itervalues()) * 0.01
     cutboxes = [(b, n)
                 for b, n
                 in sorted(boxes.iteritems(), key=lambda x: -x[1])
                 if n >= cutoff and b != BROADCAST]
     other = sum((n for n in boxes.itervalues() if n < cutoff))
     aliases = pcapdata.aliases
     if pcapdata.show_hosts:
       checked = dict((h, 1) for h in pcapdata.show_hosts)
     else:
       checked = {}
       for b, n in cutboxes:
         checked[b] = (n > cutoff * 10)
     if not pcapdata.show_fields:
       pcapdata.show_fields = DEFAULT_FIELDS
     for b in boxes.keys():
       if b not in aliases:
         aliases[b] = capdefault.aliases.get(b, b)
     self.render('view.html',
                 blob=blob_info,
                 boxes=cutboxes,
                 other=other,
                 aliases=aliases,
                 checked=checked,
                 obj=pcapdata,
                 show_fields=dict((i, 1) for i in pcapdata.show_fields),
                 avail_fields=AVAIL_FIELDS)


 class SaveHandler(_BaseHandler):

   @GoogleLoginRequired
   def post(self, blobres):
     blob_info = blobstore.BlobInfo.get(str(urllib.unquote(blobres)))
     capdefault = PcapData.GetDefault()
     u = users.get_current_user()
     if u:
       email = u.email()
     else:
       email = 'anonymous'
     sys.stderr.write('stupid user:%r email:%r\n' % (u, email))
     pcapdata = PcapData.GetOrInsertFromBlob(blob_info)
     boxes = _Boxes(blob_info)
     pcapdata.show_hosts = []
     for b in boxes.keys():
       alias = self.request.get('name-%s' % b)
       if alias:
         pcapdata.aliases[b] = alias
         capdefault.aliases[b] = alias
       else:
         pcapdata.aliases[b] = b
       if self.request.get('show-%s' % b):
         pcapdata.show_hosts.append(b)

     pcapdata.show_fields = []
     for k in AVAIL_FIELDS:
       if self.request.get('key-%s' % k):
         pcapdata.show_fields.append(k)

     _MaybeCache('on' == self.request.get('update-cache'),
                 blob_info, pcapdata)

     capdefault.put()
     pcapdata.put()

     self.redirect('/d3viz.html?key=%s&to_plot=%s'
                   % (_Esc(str(blob_info.key())),
                      _Esc(','.join(pcapdata.show_fields))))


 def _MaybeCache(update_cache, blob_info, pcapdata):
   """Update cache when asked to do so. Cache when no cache found."""

   if update_cache:
     pcapdata.js_packets = None
     pcapdata.js_streams = None

   if pcapdata.js_packets is not None:
     print "We just used cache, didn't we"
     return

   reader = blob_info.open()
   begin = time.time()

   j = []
   pairs = set()
   for i, (p, unused_frame) in enumerate(wifipacket.Packetize(reader)):
     if IS_DEBUG and i > SAMPLE_SIZE:
       print 'Done', i
       break
     j.append(p)
     pairs.add((p.get('ta', 'no_ta'), p.get('ra', 'no_ra')))

   pairs_dict = [{'ta': t[0], 'ra': t[1]} for t in pairs]

   pcapdata.js_packets = json.dumps(j)
   pcapdata.js_streams = json.dumps(pairs_dict)

   end = time.time()
   print 'Spent on caching', (end - begin), 'sec'


 class JsonHandler(_BaseHandler):

   @GoogleLoginRequired
   def get(self, blobres):
     blob_info = blobstore.BlobInfo.get(str(urllib.unquote(blobres)))
     pcapdata = PcapData.GetOrInsertFromBlob(blob_info)

     self.response.headers['Content-Type'] = 'application/json'
     js_bundle = {
         'js_packets': pcapdata.js_packets,
         'js_streams': pcapdata.js_streams,
     }
     self.response.out.write(json.dumps(js_bundle))


 def Handle500(unused_req, resp, exc):
   resp.clear()
   resp.headers['Content-type'] = 'text/plain'
   resp.write(traceback.format_exc(exc))
   resp.set_status(500)


 settings = dict(
     debug=1,
 )

 wsgi_app = webapp2.WSGIApplication([
     (r'/', MainHandler),
     (r'/upload', UploadHandler),
     (r'/download/([^/]+)/[^/]+$', DownloadHandler),
     (r'/view/([^/]+)$', ViewHandler),
     (r'/save/([^/]+)$', SaveHandler),
     (r'/json/([^/]+)$', JsonHandler),
 ], **settings)

 wsgi_app.error_handlers[500] = Handle500
	# Copyright 2014 Google Inc. All Rights Reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	"""Handlers for uploading, filtering, and visualization."""

	import collections
	import errno
	import json
	import sys
	import time
	import traceback
	import urllib
	import webapp2
	import tornado.template
	import wifipacket
	from google.appengine.api import memcache
	from google.appengine.api import users
	from google.appengine.ext import blobstore
	from google.appengine.ext import ndb
	from google.appengine.ext.webapp import blobstore_handlers

	BROADCAST = 'ff:ff:ff:ff:ff:ff'
	DEFAULT_FIELDS = ['seq', 'rate']
	AVAIL_FIELDS = ['seq', 'mcs', 'spatialstreams', 'bw', 'rate', 'retry',
	'type', 'typestr', 'dbm_antsignal', 'dbm_antnoise',
	'bad']

	IS_DEBUG = False
	SAMPLE_SIZE = 2

	loader = tornado.template.Loader('.')


	def _Esc(s):
	"""Like tornado.escape.url_escape, but only escapes &, #, %, and =."""
	out = []
	for c in s:
	if c in ['&', '#', '%', '=']:
	out.append('%%%02X' % ord(c))
	else:
	out.append(c)
	return ''.join(out)


	def AllowedEmails():
	try:
	return open('email-allow.txt').read().split()
	except IOError as e:
	if e.errno == errno.ENOENT:
	pass
	else:
	raise
	return []


	def GoogleLoginRequired(func):
	"""Enforcing @google.com login."""

	def Handler(self, args, *kwargs):
	user = users.get_current_user()
	if not user:
	self.redirect(users.create_login_url('/'))
	elif (not user.email().endswith('@google.com') and
	user.email() not in AllowedEmails()):
	self.response.set_status(401, 'Unauthorized')
	self.response.write("Sorry. You're not an authorized user.")
	else:
	return func(self, args, *kwargs)
	return Handler


	class PcapData(ndb.Model):
	"""Info about pcap file and its visualization settings/cache."""

	create_time = ndb.DateTimeProperty(auto_now_add=True)
	create_user_email = ndb.StringProperty()
	filename = ndb.StringProperty()
	show_hosts = ndb.StringProperty(repeated=True)
	show_fields = ndb.StringProperty(repeated=True)
	aliases = ndb.PickleProperty()

	# Cached JSON representations for various useful subsets of data
	# to be passed to the JS side for visualization.
	# All captured packets
	js_packets = ndb.JsonProperty(compressed=True)
	# All pairs of (transmitter, receiver)
	js_streams = ndb.JsonProperty(compressed=True)

	@staticmethod
	def GetDefault():
	return PcapData.get_or_insert(str('*'), show_hosts=[], aliases={})

	@staticmethod
	def GetOrInsertFromBlob(blob_info):
	u = users.get_current_user()
	if u:
	email = u.email()
	else:
	email = '<anonymous>'
	return PcapData.get_or_insert(str(blob_info.key()),
	show_hosts=[], aliases={},
	filename=blob_info.filename,
	create_user_email=email)


	class _BaseHandler(webapp2.RequestHandler):

	def render(self, template, **kwargs):
	d = dict()
	d.update(kwargs)
	self.response.write(loader.load(template).generate(**d))


	class MainHandler(_BaseHandler):

	@GoogleLoginRequired
	def get(self):
	upload_url = blobstore.create_upload_url('/upload')
	q = PcapData.query().order(-PcapData.create_time).fetch(10)
	self.render('index.html',
	upload_url=upload_url,
	recents=q)


	class UploadHandler(blobstore_handlers.BlobstoreUploadHandler):

	def post(self):
	upload_files = self.get_uploads()
	sys.stderr.write('upload: %r\n' % upload_files)
	blob_info = upload_files[0]
	reader = blob_info.open()
	try:
	wifipacket.Packetize(reader).next() # just check basic file header
	except wifipacket.Error:
	blob_info.delete()
	raise
	self.redirect('/view/%s' % blob_info.key())


	class DownloadHandler(blobstore_handlers.BlobstoreDownloadHandler):

	def get(self, blobres):
	blob_info = blobstore.BlobInfo.get(str(urllib.unquote(blobres)))
	self.send_blob(blob_info)


	def _Boxes(blob_info):
	"""Re-/store from/to memcache number of packets per mac address."""

	boxes = memcache.get(str(blob_info.key()), namespace='boxes')
	if not boxes:
	reader = blob_info.open()
	boxes = collections.defaultdict(lambda: 0)
	# TODO(katepek): use cache here instead if available
	for p, unused_frame in wifipacket.Packetize(reader):
	if 'flags' in p and p.flags & wifipacket.Flags.BAD_FCS: continue
	if 'ta' in p and 'ra' in p:
	boxes[p.ta] += 1
	boxes[p.ra] += 1
	memcache.add(key=str(blob_info.key()), value=dict(boxes),
	namespace='boxes')
	return boxes


	class ViewHandler(_BaseHandler):

	@GoogleLoginRequired
	def get(self, blobres):
	blob_info = blobstore.BlobInfo.get(str(urllib.unquote(blobres)))
	capdefault = PcapData.GetDefault()
	pcapdata = PcapData.GetOrInsertFromBlob(blob_info)

	boxes = _Boxes(blob_info)
	cutoff = max(boxes.itervalues()) * 0.01
	cutboxes = [(b, n)
	for b, n
	in sorted(boxes.iteritems(), key=lambda x: -x[1])
	if n >= cutoff and b != BROADCAST]
	other = sum((n for n in boxes.itervalues() if n < cutoff))
	aliases = pcapdata.aliases
	if pcapdata.show_hosts:
	checked = dict((h, 1) for h in pcapdata.show_hosts)
	else:
	checked = {}
	for b, n in cutboxes:
	checked[b] = (n > cutoff * 10)
	if not pcapdata.show_fields:
	pcapdata.show_fields = DEFAULT_FIELDS
	for b in boxes.keys():
	if b not in aliases:
	aliases[b] = capdefault.aliases.get(b, b)
	self.render('view.html',
	blob=blob_info,
	boxes=cutboxes,
	other=other,
	aliases=aliases,
	checked=checked,
	obj=pcapdata,
	show_fields=dict((i, 1) for i in pcapdata.show_fields),
	avail_fields=AVAIL_FIELDS)


	class SaveHandler(_BaseHandler):

	@GoogleLoginRequired
	def post(self, blobres):
	blob_info = blobstore.BlobInfo.get(str(urllib.unquote(blobres)))
	capdefault = PcapData.GetDefault()
	u = users.get_current_user()
	if u:
	email = u.email()
	else:
	email = 'anonymous'
	sys.stderr.write('stupid user:%r email:%r\n' % (u, email))
	pcapdata = PcapData.GetOrInsertFromBlob(blob_info)
	boxes = _Boxes(blob_info)
	pcapdata.show_hosts = []
	for b in boxes.keys():
	alias = self.request.get('name-%s' % b)
	if alias:
	pcapdata.aliases[b] = alias
	capdefault.aliases[b] = alias
	else:
	pcapdata.aliases[b] = b
	if self.request.get('show-%s' % b):
	pcapdata.show_hosts.append(b)

	pcapdata.show_fields = []
	for k in AVAIL_FIELDS:
	if self.request.get('key-%s' % k):
	pcapdata.show_fields.append(k)

	_MaybeCache('on' == self.request.get('update-cache'),
	blob_info, pcapdata)

	capdefault.put()
	pcapdata.put()

	self.redirect('/d3viz.html?key=%s&to_plot=%s'
	% (_Esc(str(blob_info.key())),
	_Esc(','.join(pcapdata.show_fields))))


	def _MaybeCache(update_cache, blob_info, pcapdata):
	"""Update cache when asked to do so. Cache when no cache found."""

	if update_cache:
	pcapdata.js_packets = None
	pcapdata.js_streams = None

	if pcapdata.js_packets is not None:
	print "We just used cache, didn't we"
	return

	reader = blob_info.open()
	begin = time.time()

	j = []
	pairs = set()
	for i, (p, unused_frame) in enumerate(wifipacket.Packetize(reader)):
	if IS_DEBUG and i > SAMPLE_SIZE:
	print 'Done', i
	break
	j.append(p)
	pairs.add((p.get('ta', 'no_ta'), p.get('ra', 'no_ra')))

	pairs_dict = [{'ta': t[0], 'ra': t[1]} for t in pairs]

	pcapdata.js_packets = json.dumps(j)
	pcapdata.js_streams = json.dumps(pairs_dict)

	end = time.time()
	print 'Spent on caching', (end - begin), 'sec'


	class JsonHandler(_BaseHandler):

	@GoogleLoginRequired
	def get(self, blobres):
	blob_info = blobstore.BlobInfo.get(str(urllib.unquote(blobres)))
	pcapdata = PcapData.GetOrInsertFromBlob(blob_info)

	self.response.headers['Content-Type'] = 'application/json'
	js_bundle = {
	'js_packets': pcapdata.js_packets,
	'js_streams': pcapdata.js_streams,
	}
	self.response.out.write(json.dumps(js_bundle))


	def Handle500(unused_req, resp, exc):
	resp.clear()
	resp.headers['Content-type'] = 'text/plain'
	resp.write(traceback.format_exc(exc))
	resp.set_status(500)


	settings = dict(
	debug=1,
	)

	wsgi_app = webapp2.WSGIApplication([
	(r'/', MainHandler),
	(r'/upload', UploadHandler),
	(r'/download/([^/]+)/[^/]+$', DownloadHandler),
	(r'/view/([^/]+)$', ViewHandler),
	(r'/save/([^/]+)$', SaveHandler),
	(r'/json/([^/]+)$', JsonHandler),
	], **settings)

	wsgi_app.error_handlers[500] = Handle500