# File from polymorph project
# Copyright (C) 2020 Santiago Hernandez Ramos <shramos@protonmail.com>
# For more information about the project: https://github.com/shramos/polymorph

from pyshark import FileCapture
from polymorph.tfield import TField
from polymorph.template import Template
from polymorph.tlayer import TLayer
from collections import OrderedDict


class TGenerator(object):
    """This class is responsible for generating the templates from tshark packets."""

    def __init__(self, pcapfile, userfilter=None):
        """Initialization method of the class.

        Parameters
        ----------
        pcapfile : str
            Path to a previously captured pcap.

        """
        self._tshark_pkts = FileCapture(
            input_file=pcapfile,
            keep_packets=False,
            use_json=True,
            include_raw=True,
            display_filter=userfilter)

    def __iter__(self):
        return self

    def __next__(self):
        """Generates a list of templates from a tshark file

        Returns
        -------
        :obj:`Template`
            A `Template` generated from the packets.

        """
        try:
            tshark_pkt = self._tshark_pkts.next()
            return self.tgenerate(tshark_pkt)
        except UnicodeDecodeError:
            self.__next__()
        except StopIteration:
            self._tshark_pkts.close()
            raise

    def tgenerate(self, tshark_pkt):
        """Generates a template from a tshark packet.

        Parameters
        ----------
        tshark_pkt : :obj:``
            The packet generated by Pyshark.

        Returns
        -------
        :obj: `Template`
            The `Tempalate` generated from the packets.

        """
        pkt_raw = tshark_pkt.get_raw_packet()
        pkt_len = len(tshark_pkt)
        # Generation of the `Template`
        template = Template(
            "Building", raw=pkt_raw, description="Template automatically generated by polymorph.")
        # Filtering and adding the `TLayer` objects to the `Template`
        tshark_layers = [
            l for l in tshark_pkt.layers if "_raw" not in l.layer_name.lower()]
        for tshark_layer in tshark_layers:
            # Generation of the `Tlayer`
            try:
                lraw = getattr(
                    tshark_pkt, tshark_layer.layer_name + "_raw").value
            except AttributeError:
                continue
            if type(lraw) is list:
                lraw = lraw[0]
            lraw = bytes.fromhex(lraw)
            tlayer_len = len(lraw)
            try:
                tlayer_start = pkt_raw.index(lraw)
                lslice = slice(tlayer_start, tlayer_start + tlayer_len)
                tlayer = TLayer(name=tshark_layer.layer_name.upper(),
                                lslice=lslice, pkt_raw=pkt_raw)
                # Filtering and adding the `TField` objects to the `TLayer`
                for f in self._get_tlayer_fields(tshark_layer, tlayer, pkt_raw):
                    tlayer.addfield(f)
            # If it corresponds to a reassembled layer
            except ValueError:
                lslice = slice(0, tlayer_len)
                lname = tshark_layer.layer_name.upper() + " (REASSEMBLED)"
                tlayer = TLayer(name=lname,
                                lslice=lslice, pkt_raw=lraw)
                # Filtering and adding the `TField` objects to the `TLayer`
                for f in self._get_tlayer_fields(tshark_layer, tlayer, lraw):
                    tlayer.addfield(f)
            # Adding the `TLayer` to the `Template`
            template.addlayer(tlayer)
        # We changed the name of the template
        template.name = self._name(template.layernames())
        return template

    def _get_tlayer_fields(self, tshark_layer, tlayer, pkt_raw):
        """Generates template fields for a given tshark layer."""
        tshark_fields = self._traverse_fields(tshark_layer._all_fields)
        for tname, tvalue in tshark_fields.items():
            # Building the TField
            tfield = TField(
                fname=tname,
                fslice=slice(tvalue[2], tvalue[2]+tvalue[3]),
                fsize=tvalue[3],
                pkt_raw=pkt_raw,
                trepr=tvalue[0],
                ttype=tvalue[5],
                tmask=tvalue[4],
                layer=tlayer)
            yield tfield

    def _traverse_fields(self, all_fields):
        """Traverse and parse all tshark fields from a tshark layer."""
        parse_fields = OrderedDict()
        # Filter all fields without _raw
        try:
            fields = [f for f in all_fields.keys(
            ) if "_raw" not in f and "_tree" not in f]
        except:
            return parse_fields
        for f in fields:
            # If not present the equivalent field with _raw it is a title
            if f + "_raw" not in all_fields.keys():
                if type(all_fields[f]) is str:
                    continue
                elif type(all_fields[f]) is list:
                    for subf in range(len(all_fields[f])):
                        newfield = self._traverse_fields(all_fields[f][subf])
                        parse_fields = self._update(parse_fields, newfield)
                    continue
                else:
                    newfield = self._traverse_fields(all_fields[f])
                    parse_fields = self._update(parse_fields, newfield)
                    continue
            # Get the _raw equivalent of the field
            if f + "_raw_raw" in all_fields.keys():
                raw_f = all_fields[f + "_raw_raw"]
            else:
                raw_f = all_fields[f + "_raw"]
            # If this condition match, f is a field
            if type(raw_f[-1]) is int and raw_f[-1] != 0:
                # Some fields can have subfields
                if type(all_fields[f]) is dict:
                    newfield = self._traverse_fields(all_fields[f])
                    parse_fields = self._update(parse_fields, newfield)
                    continue
                else:
                    new_f = f.split(".")
                    if len(new_f) == 1:
                        new_f = new_f[0]
                    elif len(new_f) > 1:
                        new_f = ".".join(new_f[1:])
                    parse_fields = self._update(
                        parse_fields, {new_f: [all_fields[f]] + raw_f})
            elif type(raw_f[-1]) is int and raw_f[-1] == 0:
                if type(all_fields[f]) is str:
                    raw_f[-1] = -1
                    newfield = self._traverse_fields({f + "_raw": raw_f,
                                                      f: all_fields[f]})
                    parse_fields = self._update(parse_fields, newfield)
                else:
                    newfield = self._traverse_fields(all_fields[f])
                    parse_fields = self._update(parse_fields, newfield)
            elif type(raw_f[-1]) is list:
                for subf in range(len(raw_f)):
                    newfield = self._traverse_fields({f + "_raw": raw_f[subf],
                                                      f: all_fields[f][subf]})
                    parse_fields = self._update(parse_fields, newfield)
            else:
                print("[ERROR] Parsing field:", f)
        # Return the parsed fields without duplicates
        return parse_fields

    def _update(self, od1, od2):
        """Updates the list of fields without duplicate names."""
        for k, v in od2.items():
            # Delete duplicate fields
            duplicate = False
            for k2, v2 in od1.items():
                if v[2] == v2[2] and v[3] == v2[3] and v[4] == v2[4]:
                    duplicate = True
            if duplicate:
                continue
            new_k = self._new_fname(k, od1)
            od1[new_k] = v
        return od1

    def _new_fname(self, fname, all_fields):
        """Generates new names for fields with duplicate name."""
        if fname in all_fields:
            if fname[-1].isnumeric():
                return self._new_fname(fname[:-1] + str(int(fname[-1]) + 1), all_fields)
            else:
                return self._new_fname(fname + str(2), all_fields)
        else:
            return fname

    def _name(self, layers):
        """Generates a name for each template based on the layers of the
        packet."""
        return "Template: " + " / ".join(layers)
