#!/usr/bin/python
# -*- coding: utf-8 -*-
import sys
import subprocess
import re
import socket
import datetime
from datetime import date
import os, errno
import argparse
import tempfile
from lxml import etree
from xml.dom.minidom import parse, parseString

#####################
# program arguments
#####################
parser = argparse.ArgumentParser()
parser.add_argument("--pcap_file", "-p", type=str,help="input pcap file to be translated")
args = parser.parse_args()

#####################
# get xml document from pcap
#####################
orig_pcap_file_name = args.pcap_file.strip()
orig_pdml_string = subprocess.check_output(["tshark", '-T', 'pdml', '-r', orig_pcap_file_name])

orig_dom = parseString(orig_pdml_string)

#####################
# filtering unwanted packets
#####################
#cases = orig_etree.findall(".//proto[@name='sctp']")
packets = orig_dom.getElementsByTagName("packet")
for packet in packets:
    found_sctp = False
    found_s1ap = False
    sctp_node = None
    frame_node = None
    ip_node = None
    protos = packet.getElementsByTagName("proto")
    for proto in protos:
        attrs   = proto.attributes
        urlnode = attrs['name']
        if urlnode.nodeValue == 'frame':
            frame_node = proto
        elif urlnode.nodeValue == 'ip':
            frame_node.appendChild(proto)
            ip_node = proto
        elif urlnode.nodeValue == 'sctp':
            found_sctp = True
            ip_node.appendChild(proto)
            sctp_node = proto
        elif urlnode.nodeValue == 's1ap':
            found_s1ap = True
            sctp_node.appendChild(proto)
        elif urlnode.nodeValue == 'geninfo':
            packet.removeChild(proto)
        elif urlnode.nodeValue == 'eth':
            packet.removeChild(proto)
    if found_sctp == False:
        # hopefully it seems to work (remove iterated packet)
        packet.parentNode.removeChild(packet)

#####################
# dom to xml string
#####################
filtered_pdml_string = orig_dom.toprettyxml(indent="  ")
cleaned_pdml_string = ""
#####################
# remove blank lines in xml string
#####################
lines = filtered_pdml_string.splitlines()
for line in lines:
    if len(line.strip()):
        if line[:-1]:
            cleaned_pdml_string += line + '\r\n'
#print "'%s'" %  cleaned_pdml_string 
#####################
# write pdml string to pdml file
#####################
out_pdml_file_name = os.path.dirname(orig_pcap_file_name) + '/' + os.path.splitext(os.path.basename(orig_pcap_file_name))[0] + '.pdml'
out_file = open(out_pdml_file_name, "w")
out_file.write(cleaned_pdml_string)
out_file.close()

############################################################
# DECEIVING HTML BONUS: DO NOT SEEM TO WORK CORRECTLY IN FIREFOX
# DID NOT INVESTIGATE
#####################
# write xml string to html file
#####################
xsl_root = etree.fromstring(open('/usr/share/wireshark/pdml2html.xsl').read())
transform = etree.XSLT(xsl_root)
xml_root = etree.fromstring(cleaned_pdml_string)
trans_root = transform(xml_root)
filtered_html_string = etree.tostring(trans_root)
#####################
# write html string to html file
#####################
out_html_file_name = os.path.dirname(orig_pcap_file_name) + '/' +os.path.splitext(os.path.basename(orig_pcap_file_name))[0] + '.html'
out_file = open(out_html_file_name, "w")
out_file.write(filtered_html_string)
out_file.close()