#!/usr/local/bin/python

# Send/receive UDP multicast heartbeat packets
#
# Usage:
#   heartbeat [OPTIONS] 
#
# Examples:
#   # generate heartbeats every 1000ms describing system and neighbors on wlan0
#   # that are active within past 5000ms
#   heartbeat -i wlan0 -T 5000 -P 1000
#
# To get current aggregated node list: 
#   telnet localhost 8888
#
# Notes:
#   - must have a viable route for mcast (either a default route or explicit):
#     route add default gw 172.20.0.31 wlan0
#    or
#     route add -net 224.0.0.0 netmask 224.0.0.0 wlan0
#   - make sure to have getip installed
#
# TODO:
#   - skip packets that don't match our protocol major version
#     (major version bumps break compat, minor version bumps introduce compat features)
#

GROUP='225.0.0.250'
HOST='127.0.0.1'
PORT=8888
HBPORT=8123
TTL=1
PERIOD=1
WIFI='wlan3'
INACTIVE=0
PKT_VER="1.0"
VERBOSE=0
BATMAN=False

import sys
import getopt
import time
import struct
import string
import os
import fcntl
import shlex
import subprocess 
import re
from threading import Thread
from socket import *
#from collections import namedtuple
import collections 

OCU_IP=os.environ['MONIP']
WIFI=os.environ['WLANMESH']

def verbose(lvl, msg):
   if (lvl <= VERBOSE):
      print msg

FILTER=''.join([(len(repr(chr(x)))==3) and chr(x) or '.' for x in range(256)])
def hexdump(src, length=8):
    N=0; result=''
    while src:
       s,src = src[:length],src[length:]
       hexa = ' '.join(["%02X"%ord(x) for x in s])
       s = s.translate(FILTER)
       result += "%04X   %-*s   %s\n" % (N, length*3, hexa, s)
       N+=length
    return result


# Encapsulates a heartbeat message
PKT_FMT = '!BBHL6BLB' # pkt_vermaj, pkt_vermin, pkt_size, uptime, Mac_ADDR, IP_ADDR, num_stations                                      
STA_FMT = '!H6BLLLLLLLhBf' # size, mac, ip_addr, metric, inactive, rx_bytes, rx_packets, tx_bytes, tx_packets, signal, tx_rate, throughput   
class heartbeat:
   # unpack data to simple string representation
   #   TODO: perhaps use simplejson: http://simplejson.googlecode.com/svn/tags/simplejson-2.1.0/docs/index.html
   #    or python-json, or python-cjson
   @staticmethod
   def str(data, arp = None, mac = None):
      s = ''
      i = struct.calcsize(PKT_FMT)
      (pkt_vermaj, pkt_vermin, pkt_size, uptime, m0, m1, m2, m3, m4, m5, ip_addr, num_stations) = struct.unpack(PKT_FMT, data[0:i])
      s += "  uptime:%d\n" % uptime
      s += "  mac:%02x:%02x:%02x:%02x:%02x:%02x" % (m0, m1, m2, m3, m4, m5)
      s += "  ip_addr:%s\n" % ip_addr
      s += "  stations:%d\n" % num_stations
      for j in range(num_stations):
         (size, m0, m1, m2, m3, m4, m5, ip_addr, metric, inactive, rx_bytes, rx_packets, tx_bytes, tx_packets, signal, tx_rate, throughput) = struct.unpack_from(STA_FMT, data, i)
         mac = "%02x:%02x:%02x:%02x:%02x:%02x" % (m0, m1, m2, m3, m4, m5)
         i += struct.calcsize(STA_FMT)
         sta  = "    station%d:\n" % (j+1)
         sta += "      ip:%s\n" % ip_addr
         sta += "      metric:%s\n" % metric 
         sta += "      mac:%s\n" % mac
         sta += "      inactive:%d\n" % inactive
         sta += "      rx_bytes:%d\n" % rx_bytes
         sta += "      rx_packets:%d\n" % rx_packets
         sta += "      tx_bytes:%d\n" % tx_bytes
         sta += "      tx_packets:%d\n" % tx_packets
         sta += "      signal:%d\n" % signal
         sta += "      tx_rate:%d\n" % tx_rate
         sta += "      throughput:%f\n" % throughput
         #print sta
         s += sta
      return s
      

   # generate and return a heartbeat packet
   #  - uses /sys/kernel/debug/ieee80211/<phy>/stations as this is 
   #    probably a more stable API than output of 'iw dev <iface> station dump'
   @staticmethod
   def create(iface):
      data = ''
      phy = "phy%s" % open("/sys/class/net/%s/phy80211/index" % iface).read().strip()
      uptime, idletime = [float(f) for f in open("/proc/uptime").read().split()]
      # station data
      i = 0
      stations = []
      src_mac = [0, 0, 0, 0, 0, 0]
      src_mac_file = open("/sys/class/net/%s/address" % iface).read().strip()
      f = src_mac_file.split(':')
      for j in range(len(f)):
         src_mac[j] = int(f[j], 16)
      

      dir = "/sys/kernel/debug/ieee80211/%(phy)s/netdev:%(iface)s/stations/" % \
	{'phy':(phy), 'iface': (iface)}
      for mac in os.listdir(dir):
         inactive_ms = int(open(dir + mac + '/inactive_ms').read());
         if INACTIVE == 0 or inactive_ms < INACTIVE:
            i += 1
            #verbose(1, "   station%d:%s:%d" % (i, mac, inactive_ms))
            tx_rate = 0
            throughput = 0.0
            try:
               for line in open(dir + mac + '/rc_stats').read().split('\n'):
                  if len(line) and line[0] == 'T':
                     tx_rate = int(float(line[3:8]))
                     throughput = float(line[11:18])
                     break
            except:
               pass
            m = [0,0,0,0,0,0]
            b = mac.split(':')
            for j in range(len(b)):
               m[j] = int(b[j], 16)
            ip_addr = 0
            metric = 0
            sta_data = struct.pack(STA_FMT, 
               struct.calcsize(STA_FMT),
               m[0], m[1], m[2], m[3], m[4], m[5],
               ip_addr,
               metric,
               inactive_ms,
               int(open(dir + mac + '/rx_bytes').read()),
               int(open(dir + mac + '/rx_packets').read()),
               int(open(dir + mac + '/tx_bytes').read()),
               int(open(dir + mac + '/tx_packets').read()),
               int(open(dir + mac + '/last_signal').read()),
               tx_rate,
               throughput,
            )
            data += sta_data

      
      # get ip address of this interface
      if BATMAN:
            args = shlex.split("getip bat0")
            print "BATMAN!!!"
      else:
            args = shlex.split("getip %s" % iface)
            print "NOT BATMAN!!!"
      print args
      p = subprocess.Popen(args, stdout=subprocess.PIPE)
      r = p.communicate()[0]
      p.wait()
      print >>sys.stderr, r
      
      #TAKE THIS OUT LATER!!!
      r = os.environ['MESHIP']
      # heartbeat packet
      data = struct.pack(PKT_FMT, 
            int(PKT_VER.split('.')[0]),
            int(PKT_VER.split('.')[1]),
            struct.calcsize(PKT_FMT) + i*struct.calcsize(STA_FMT), 
            int(uptime),
            src_mac[0], src_mac[1], src_mac[2], src_mac[3], src_mac[4], src_mac[5],
            struct.unpack('I', inet_aton(r))[0],
            i) + data
      return data


# Send periodic heartbeat packets
class sender(Thread):

   def __init__(self, iface, group = '<broadcast>', port = HBPORT, period = PERIOD):
      Thread.__init__(self, name='Sender')
      self.iface = iface
      self.port = port;
      self.group = group;
      self.period = period
      self.done = False

   def stop(self):
      self.done = 1;

   def run(self):
      verbose(0, "Monitoring nodes from %s" % self.iface)
      s = socket(AF_INET, SOCK_DGRAM)
      '''
      if self.group == '<broadcast>':
         s.setsockopt(SOL_SOCKET, SO_BROADCAST, 1)
         verbose(0, "Sending broadcast heartbeats")
      else:
         ttl = struct.pack('b', TTL)
         s.setsockopt(IPPROTO_IP, IP_MULTICAST_TTL, ttl)
         verbose(0, "Sending multicast heartbeats to %s:%d TTL=%d" % (self.group, self.port, TTL))
      '''
      while not self.done:
         data = heartbeat.create(self.iface)
         verbose(4, "sending %d byte heartbeat" % len(data))
         verbose(5, heartbeat.str(data))
         s.sendto(data, (self.group, self.port))
         time.sleep(self.period)
      verbose(0, "sender complete");


def main():
   global GROUP
   global HOST
   global PORT
   global HBPORT
   global TTL
   global PERIOD
   global WIFI
   global INACTIVE
   global VERBOSE
   global OCU_IP
   global BATMAN

   def usage():
      print "usage: %s [OPTIONS]" % sys.argv[0]
      print ""
      print "  Options:"
      print "    -h,--help      - This usage"
      print "    -v,--verbose   - Set verbosity level"
      print ""
      print "  Sender Options:"
      print "    -o,--ocu       - IP address of OCU on mesh (default=%s)" % OCU_IP
      print "    -P,--period    - Period between sending (defalt=%dsecond)" % PERIOD 
      print "    --hbport       - Port to send/recv heartbeats on (default=%d)" % HBPORT
      print "    -i,--interface - Wireless Interface to monitor neighbors of (default=%s)" % WIFI
      print "    -t,--ttl       - Multicast TTL (default=%d)" % TTL
      print "    -T,--timeout   - Timeout nodes that have active_tx longer than this (0=disabled) (default=%d)" % INACTIVE
      print "    -B,--batman    - Use heartbeat over batman (default=%s)" % BATMAN
      print ""
      print "  Receiver Options:"
      print "    -H,--host      - Host to bind to on receive (default=%s)" % HOST
      print "    -p,--port      - Port to listen on (default=%d)" % PORT
      print ""

   try:
      opts, args = getopt.getopt(sys.argv[1:], "hbp:o:P:H:g:t:i:T:v:B", 
         ["help", "broadcast", "period=", "hbport=", "port=", "host=",
          "group=", "ttl=", "interface=", "timeout=", "verbose=","ocu=","batman"]);
   except getopt.GetoptError, err:
      print str(err)
      usage()
      sys.exit(2)
   send = OCU_IP
   for o, a in opts:
      if o in ("-h", "--help"):
         usage()
         sys.exit()
      elif o in ("-o", "--ocu"):
         OCU_IP = o
      elif o in ("-P", "--period"):
         PERIOD = int(a)
      elif o in ("--hbport"):
         HBPORT = int(a)
      elif o in ("-p", "--port"):
         PORT = int(a)
      elif o in ("-H", "--host"):
         HOST = a
      elif o in ("-t", "--ttl"):
         TTL = int(a)
      elif o in ("-i", "--interface"):
         WIFI = a
      elif o in ("-T", "--timeout"):
         INACTIVE = int(a)
      elif o == "-v":
         VERBOSE = VERBOSE + 1
      elif o == "--verbose":
         VERBOSE = int(a)
      elif o in ("-B", "--batman"):
            BATMAN = True
            print "YEAH! BATMAN!"

   verbose(0, "packet_size = %d + (stations * %d)" %
      (struct.calcsize(PKT_FMT), struct.calcsize(STA_FMT)))


   print "send is %s" % send
   s = sender(WIFI, OCU_IP, HBPORT)
   s.start()

   # listen for client connection
   verbose(0, "Listening on %s:%d for connections" % (HOST, PORT))
   sock = socket(AF_INET, SOCK_STREAM)
   sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
   sock.bind((HOST, PORT))
   sock.listen(1)
   try:
      while 1:
         conn, addr = sock.accept()
         verbose(1, "Servicing incomming connection from %s:%d" % addr)
         # handle commands
         #while self.done == 0:
         #   data = conn.recv(1024)
         #   if not data: break
         #   conn.send(data)

         conn.send("rx_packets:%d\n" % r.rx_packets)
         conn.send("rx_bytes:%d\n" % r.rx_bytes)
         conn.send("rx_bitrate:%dbps\n" % (r.rx_bytes*8 / (time.time()-r.rx_time) ))
         conn.send("rx_pktrate:%2.1fHz\n" % (r.rx_packets / (time.time()-r.rx_time) ))
         conn.send("nodes:%d\n" % len(r.nodes))
         # dump node list
         i = 0
         for k, v in r.nodes.iteritems():
            inactive_ms = time.time() - r.nodes_rxtime[k]
            if INACTIVE == 0 or inactive_ms < INACTIVE/1000.0:
               i = i + 1
               conn.send("node%d:%s:\n%s" % (i, k, v))
               #conn.send("node%d:\n" % i)
               #conn.send("  ip:%s\n%s" % (k, v))
         conn.close()
   except KeyboardInterrupt:
      pass
   sock.close()
   s.stop()
   s.join()

if __name__ == "__main__":
   main()