#!/usr/bin/python

# -----------------------------------------------------------------------------
#  This program takes as input a log file dumped by aimSniff. It parses the
#  file and builds a connected graph depicting the interrelation of the chat
#  participants. It outputs a text file that can be read by GraphVis.
#
#               (c) 2004 Jason Bittel
# -----------------------------------------------------------------------------

import os
import string
import sys
import re

# -----------------------------------------------------------------------------
# Nodes of the graph
# -----------------------------------------------------------------------------
class Vertex:
   def __init__(self, id):
      self.id = id

   def getID(self):
      return self.id

# -----------------------------------------------------------------------------
# Edges of the graph
# -----------------------------------------------------------------------------
class Edge:
   def __init__(self, id, u, v, count=1):
      self.id = id
      self.u = u
      self.v = v
      self.count = count

   def getID(self):
      return self.id

   def getCount(self):
      return self.count

   def getU(self):
      return self.u

   def getV(self):
      return self.v

   def incCount(self):
      self.count += 1

# -----------------------------------------------------------------------------
# Read input file and parse each line
# -----------------------------------------------------------------------------
def readFile(infile):
   "Read input file and parse each line"

   try:
      input = open(infile)

      notAlphaNum = re.compile('\W')

      for line in input:
         recFields = string.splitfields(line, "##")
         recType = recFields[recFields.index("TYPE") + 1]

         if recType == "Incoming Message" or recType == "Outgoing Message":
            source = recFields[recFields.index("FROM") + 1]
            dest = recFields[recFields.index("DESTHANDLE") + 1]

            # Skip record if either name is blank
            if source == "" or dest == "":
               continue

            # Lowercase screen names
            source = source.lower()
            dest = dest.lower()

            # Clean screen names
            source = notAlphaNum.sub('_', source)
            dest = notAlphaNum.sub('_', dest)

            # Add vertices/edge to graph
            buildGraph(source, dest)

      input.close()
   except IOError, err:
      sys.stderr.write("Error: unable to access %s: %s\n" % (infile, err.strerror))
      sys.exit()

# -----------------------------------------------------------------------------
# Check if vertex exists with given id
# -----------------------------------------------------------------------------
def vertexExists(id):
   "Check to see if a given vertex exists"

   i = 0
   for vertex in allVertices:
      if vertex.getID() == id:
         return i
      i += 1

   return -1

# -----------------------------------------------------------------------------
# Check if edge exists with given id
# -----------------------------------------------------------------------------
def edgeExists(id):
   "Check to see if a given edge exists"

   i = 0
   for edge in allEdges:
      if edge.getID() == id:
         return i
      i += 1

   return -1

# -----------------------------------------------------------------------------
# Build internal graph structure from input
# -----------------------------------------------------------------------------
def buildGraph(source, dest):
   "Build graph where nodes are participants and edges are conversations"

#   if not vertexExists(source):
#      allVertices.append(Vertex(source))
#
#   if not vertexExists(dest):
#      allVertices.append(Vertex(dest))

   # Create new edge or increment existing edge
   edgeID = source + "+" + dest
   edge = edgeExists(edgeID)
   if edge < 0:
      allEdges.append(Edge(edgeID, source, dest))
   else:
      allEdges[edge].incCount()

# -----------------------------------------------------------------------------
# Output the entire graph in GraphVis format
# -----------------------------------------------------------------------------
def printGraph(outfile):
   "Output the graph information in GraphVis format"

   try:
      outfile = open(outfile, "w")
      outfile.write("digraph G {\n")
      outfile.write("\tgraph [overlap=scale];\n")
      outfile.write("\tnode [shape=circle, style=filled, fontsize=12, height=.25];\n")
      outfile.write("\tedge [fontsize=12, arrowhead=none, arrowtail=none];\n")

      for edge in allEdges:
         outfile.write("\t" + edge.getU() + " -> " + edge.getV() + " [label=\"" + str(edge.getCount()) + "\"];\n")

      outfile.write("}\n")
      outfile.close()
   except IOError, err:
      sys.stderr.write("Error: unable to access %s: %s\n" % (outfile, err.strerror))
      sys.exit()

# =============================================================================
# Main program
# =============================================================================

allEdges = []
#allVertices = []

# Gather command line parameters
infile = "aimsniff.in"

# Read/parse input file
readFile(infile)

# Output graph text file
printGraph("aimsniff.out")

# TODO:
# * Gather commandline parameters + error check
# * Allow multiple input files
# * Generate output filename
# * Work on the freaking speed issues
#     - Eliminate some of the loops
#     - Do whatever else it takes
# * Reimplement vertices, find a need for them