#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys
import subprocess
import re
import os

def usage():
  sys.stderr.write("%s: <.pcap> [tshark additional options] 'wireshark display filter'\n" % sys.argv[0])
  sys.exit(1)

n = len(sys.argv)
if n < 3:
  usage()

capture = sys.argv[1]
dfilter = sys.argv[-1]

tshark_cmd = [ '/usr/sbin/tshark', '-n' ]
tshark_cmd.extend(sys.argv[2:-1])
tshark_cmd.append('-r')
tshark_cmd.append(sys.argv[1])
tshark_cmd.append(sys.argv[-1])

# start tshark subprocess and prepare a pipe to which it will write stdout
shark = subprocess.Popen(tshark_cmd, stdout=subprocess.PIPE)
sharkout = shark.stdout

# list of messages displayed by tshark
messages = []


while True:
  line = sharkout.readline().decode('utf-8')
  # eof encountered
  if len(line) == 0:
    break

  regex = re.compile('^ *(\d+) +(\d+\.\d+) +(\d+\.\d+\.\d+\.\d+) (→|->) (\d+\.\d+\.\d+\.\d+) (.*?)$')
  ret = regex.match(line)

  if ret != None:
    msg = {}
    msg['num'] = ret.group(1)
    msg['date'] = ret.group(2)
    msg['src'] = ret.group(3)
    msg['dst'] = ret.group(5)

    m = re.sub(r'SIP.*(?:Status|Request):\s+(.*)', r'\1', ret.group(6))
    sip_message, count = re.subn(r'INVITE\s+(\S+)\s+\| , with session description', r'INVITE SDP \1', m)
    if count == 0:
      sip_message, count = re.subn(r'200 OK\s+\| , with session description', '200 OK SDP', m)
      if count == 0:
        sip_message, count = re.subn(r'(183|180)\s+(.*?)\s+\| , with session description', r'\1 \2 SDP', m)
        if count == 0:
          sip_message, count = re.subn(r'(.*)\s+\|.*$', r'\1', m)

    #print >>sys.stderr, "m '%s'" % sip_message

    msg['msg'] = sip_message
    messages.append(msg)
  else:
    sys.stderr.write("Line [%s] not handled by regex!\n" % line)
    break


# synchronously wait for tshark termination
shark.wait()
if shark.returncode != 0:
  sys.stderr.write("tshark returned error code %d\n" % shark.returncode)
  sys.exit(1)

# list of entity
# contains IP addresses used IP datagrams exchanged in this capture
entities = []
for msg in messages:
  if msg['src'] not in entities:
    entities.append(msg['src'])
  if msg['dst'] not in entities:
    entities.append(msg['dst'])

if len(entities) == 0:
  sys.exit(1)

cp_info = '/var/log/m2/pcap_tshark/pcap_tmp/cp_info.txt'
cp_info_h = {}

if os.path.exists(cp_info):
  with open(cp_info, 'r') as f:
    for l in f:
        line = l.rstrip()
        if line:
            k, ip, who = line.split(',', 2)
            cp_info_h[k] = [who,ip]

#system_ips = []
#for key in cp_info_h:
#  if cp_info_h[key][1] == "system":
#    system_ips.append(key)

# print msc generated file on stdout
print("msc {")

# dots are not allowed in entity grammar (see mscgen grammar)
# thus, name IP address by u%d, where %d is replaced by their index in the list
line_first = ''
for i in range(0, len(entities)):
  #print "'%s' '%s' '%s' '%s'" % (entities[i], k, who, ip)
  if entities[i] in cp_info_h:
    cp_info = cp_info_h[entities[i]]
    #print cp_info
    if cp_info[0] == "op":
      line_first += 'u%d[label=\"%s\\n%s\", textbgcolor=\"#f2f2f2\"]' % (i,cp_info[1][0:15], entities[i])
    elif cp_info[0] == "tp":
      line_first += 'u%d[label=\"%s\\n%s\", textbgcolor=black, textcolor=white]' % (i,cp_info[1][0:15], entities[i])
    elif cp_info[0] == "system":
        line_first += 'u%d[label=\"%s\\n%s\", textbgcolor=silver]' % (i,cp_info[1][0:15],entities[i])
    else:
      line_first += 'u%d[label=\"%s\\n%s\", textbgcolor=silver]' % (i,cp_info[1][0:15],entities[i])
  else:
      line_first += 'u%d[label=\"%s\", textbgcolor=silver]' % (i,entities[i])

  if i < len(entities)-1:
    line_first += ','
print("  %s;" % line_first)

#line_second = ''
#for i in range(0, len(entities)):
#  if entities[i] in cp_info_h:
#    line_second += 'u%d box u%d [label=\"%s\"]' % (i,i,cp_info_h[entities[i]])
#  else:
#    line_second += 'u%d box u%d [label=\"%s\"]' % (i,i,entities[i])
#  if i < len(entities)-1:
#    line_second += ','
#print("  %s;" % line_second)

# add messages
# a message is an arrow between src and dst (IP addresses)
# and a label which is the line used by tshark to describe packet content
for msg in messages:
  src = entities.index(msg['src'])
  dst = entities.index(msg['dst'])
  if src < dst:
    #print("  ->u%d [ label = \"%s\"] ;" % (src,  msg['date']))
    if re.search(r"(Status:\s+)?[45][0-9][0-9]\s+", msg['msg']):
      print("  u%d=>u%d [ label = \"%s\", textcolour=\"red\", linecolour=\"red\"] ;" % (src, dst, msg['msg']))
    elif re.search(r"(Request:\s+)?INVITE\s+", msg['msg']):
      print("  u%d=>u%d [ label = \"%s\", textcolour=\"#73e600\", linecolour=\"#73e600\" ] ;" % (src, dst, msg['msg']))
    elif re.search(r"(Request:\s+)?(BYE|CANCEL)\s+", msg['msg']):
      print("  u%d=>u%d [ label = \"%s\", textcolour=\"blue\", linecolour=\"blue\" ] ;" % (src, dst, msg['msg']))
    else:
      print("  u%d=>u%d [ label = \"%s\" ] ;" % (src, dst, msg['msg']))
     #print("  u%d=>u%d [ label = \"%s\" ] ;" % (src, dst, msg['msg']))
  else:
    if re.search(r"(Status:\s+)?[45][0-9][0-9]\s+", msg['msg']):
      print("  u%d<=u%d [ label = \"%s\", textcolour=\"red\", linecolour=\"red\"] ;" % (dst, src, msg['msg']))
    elif re.search(r"(Request:\s+)?INVITE\s+", msg['msg']):
      print("  u%d<=u%d [ label = \"%s\", textcolour=\"#73e600\", linecolour=\"#73e600\" ] ;" % (dst, src, msg['msg']))
    elif re.search(r"(Request:\s+)?(BYE|CANCEL)\s+", msg['msg']):
      print("  u%d<=u%d [ label = \"%s\", textcolour=\"blue\", linecolour=\"blue\"] ;" % (dst, src, msg['msg']))
    else:
      print("  u%d<=u%d [ label = \"%s\" ] ;" % (dst, src, msg['msg']))
    #print("  u%d<=u%d [ label = \"%s\" ] ;" % (dst, src, msg['msg']))

print("}")