#!/usr/bin/python

import sys
from threading import Thread
import socket
import bson
import shlex
import subprocess


LISTEN_IP="0.0.0.0" # listen on any ip
LISTEN_PORT=8124
IFACE='wlan0'

RESPONSE_IP="127.0.0.1"
RESPONSE_PORT=8125

class CNCType:
   """
   Defines macros for CNC
   """
   TRACE = 'traceroute'
   PING = 'ping'

def debugOut(str):
    """
    Displays debug information
    """
    DEBUG = True
    if DEBUG:
        print str

def sendCCResponse(dict, send_ip, send_port):
    """
    Sends dict to send_ip using send_port.
    Sends a CC Response.
    """
    c = bson.dumps(dict)
    b = bson.loads(c)

    debugOut('\nsending to %s:%s' % (send_ip, send_port))
    debugOut('packet: %s' % b)

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.sendto(c,(send_ip, send_port))

def performTraceRoute( end_ip):
   # need to add code to actually do a trace route, you get the picture though
   """
   Performs a traceroute on end_ip and returns the result.
   """
   return traceips(end_ip)

def performPing(start_ip, end_ip):
   """
   Calculates average ping from start_ip to end_ip and returns the result.
   """
   # need to add code to actually do a ping, you get the picture though
   return {"ping": "complete"}

def parseCNC(packet_string, send_addr=RESPONSE_IP, send_port=RESPONSE_PORT):
   """
   Attemps to parse the CNC Packet. If succesful, runs the CNCResponder, otherwise returns debug information.
   """
   packet_data = bson.loads(packet_string)
   try:
      s = packet_data['cmd']
      print 'parseCNC %s' % s
      responder = CNCResponderThread(packet_data, send_addr, send_port)
      responder.run()
   except KeyError:
      print '\nInvalid CNC Packet recieved (No cmd key in dict): \n\t%s' % packet_data

#trace route
def traceips(dest):
   """
   Calculates the traceroute on dest, and returns the result.
   """
   args = ["traceroute", dest, "-4", "-F", "-m", "50", "-n"]
   #use -i when in larger interface
   p = subprocess.Popen(args, stdout=subprocess.PIPE)
   r = p.communicate()[0]
   p.wait()
   
   a = []
   for line in r.split('\n')[1:]:
      piece = line.split();
      if len(piece) > 0:
         a.append(piece[1])
   return a

# listen for CNC packets
class CNCListenerThread(Thread):
   """
   A Listener that appropriately listens to CNC packets.
   """

   def __init__(self, iface=IFACE, ip = LISTEN_IP, port = LISTEN_PORT):
      """
      Initializes the Listener, and allows for varying interfaces, ips, and ports.
      """
      Thread.__init__(self, name='CNCListener')
      self.iface = iface
      self.ip = ip
      self.port = port
      self.done = False

   def stop(self):
      """
      Stops the Listener.
      """
      self.done = 1;

   def run(self):
      """
      Runs the Listener.
      """
      debugOut("Listening to CNC Packets from %s %s:%s" % (self.iface, self.ip, self.port) )
      sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
      sock.setblocking(0)
      sock.bind((self.ip, self.port))

      while not self.done:
         try:
            data = sock.recv(1024)
            print "received message:", bson.loads(data)
            parseCNC(data)
         except socket.error:
            pass
         except KeyboardInterrupt:
            self.stop()


#
# Thread created when we need to respond to CNCRequest
class CNCResponderThread(Thread):
   """
   A Responder that will appropriately respond to CCRequests.
   """

   def __init__(self, data, ip = RESPONSE_IP, port = RESPONSE_PORT):
      """
      Initializes the Responder, and allows for varying interfaces, ips, and ports.
      """
      Thread.__init__(self, name='CNCResponder')
      self.ip = ip
      self.port = port
      self.data = data

   def run(self):
      """
      Runs the Responder.
      """
      packet_type = self.data['cmd']
      if ( packet_type == CNCType.PING):
         debugOut("performing ping and doing stuff")
         sendCCResponse({'cmd': 'ping', 'id':2, "data": {"latency":"20"}}, self.ip, self.port)

      elif ( packet_type == CNCType.TRACE):
         debugOut("performing traceroute and doing stuff")
         args = self.data['args']
         hops = performTraceRoute(args['dest'])
         hops.insert(0, self.ip)
         sendCCResponse({'cmd': 'traceroute', 'id': self.data['id'], "data" : {"hops": hops}}, self.ip, self.port)


def main():
   """
   Begins the receiver, and stays running until killed.
   """
   reciever = CNCListenerThread()
   reciever.start()
   
   #stay in a while loop until killed
   try:
      while 1:
         pass
   except KeyboardInterrupt:
      pass
   
   reciever.stop()

if __name__ == '__main__':
    main()