#!/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 xml.dom.minidom import parse, parseString
#from xml.etree.ElementTree import ElementTree

#####################
# 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)
#orig_etree = ElementTree()
#orig_etree.parse(orig_pcap_pdml_file.name)

#####################
# 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
    protos = packet.getElementsByTagName("proto")
    for proto in protos:
        attrs   = proto.attributes
        urlnode = attrs['name']
        if urlnode.nodeValue == 'sctp':
            found_sctp = True
        elif urlnode.nodeValue == 's1ap':
            found_s1ap = True
        elif urlnode.nodeValue == 'geninfo':
            packet.removeChild(proto)
        elif urlnode.nodeValue == 'eth':
            packet.removeChild(proto)
    if found_s1ap == False:
        # hopefully it seems to work (remove iterated packet)
        packet.parentNode.removeChild(packet)

#####################
# dom to xml string
#####################
filtered_pdml_string = orig_dom.toxml()
cleaned_pdml_string = ""
#####################
# remove blank lines in xml string
#####################
lines = filtered_pdml_string.splitlines()
for line in lines:
    if line[:-1]:
        cleaned_pdml_string += line + '\r\n'
#print "'%s'" %  cleaned_pdml_string 
#####################
# write xml 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()

