#
# Copyright(C) 2005 INL
# Written by Jean Gillaux
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
#  This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.
#
#


from IPy import IP
from nupyf.net_defs import ip_protocols_rev


class GenNat(object):
    def __init__(self, src, dst, proto, to, id=-1, sport='0', dport='0', rport=None):
        self._src = self._dst = self.to = ''
        self._proto = proto
        self._name = ''
        self.id = id
        self.icmptype=''
        self._src = None
        self._dst = None
        self._to = None
        self._iface = None #iface for rule generation
        try:
            if src:
                self._src = IP(src)
            if dst:
                self._dst = IP(dst)
            if to:
                self._to = IP(to)
        except ValueError, e:
            raise e
        # Source and destination ports for matching
        self._dport = {'l':-1, 'h':-1}
        self._sport = {'l':-1, 'h':-1}
        self._rport = {'l':-1, 'h':-1}
        for p,sp in [(sport, self._sport), (dport,self._dport), (rport, self._rport)]:
            if not p:
                continue
            if p.find(':') == -1:
                sp['h'] = sp['l'] = int(p)
            else:
                l = p.split(':')
                sp['l'] = l[0]
                sp['h'] = l[1]

    def __str__(self):
        s = "ID: %i, src: %s, dst: %s, to: %s" %(self.id, repr(self._src), repr(self._dst), repr(self._to))
        s += ", sport: %s, dport: %s" %(repr(self._sport), repr(self._dport))
        return s

    def port_defined(self,port):
        return (port != {'l':-1, 'h':-1})

    def proto_defined(self):
        if self._proto:
            return True
        return False

    def set_name(self, name):
        self._name =  name

    def set_iface(self, iface):
        self._iface = iface

    def get_iface(self):
        return self._iface

class Snat(GenNat):
    def __init__(self, src, dst, proto, to, id=-1, sport='1024:65536', dport='0'):
        super(Snat,self).__init__(src, dst, proto, to, id=id, sport=sport, dport=dport)


class Dnat(GenNat):
    def __init__(self, src, dst, proto, to, id=-1, sport='1024:65536', dport='0',rport=None):
        super(Dnat, self).__init__(src, dst, proto, to, id=id, sport=sport, dport=dport,rport=rport)


class Pnat(GenNat):
    def __init__(self, src, dst, proto, to, id=-1, sport='1024:65536', dport='0',rport=None):
        super(Pnat, self).__init__(src, dst, proto, to, id=id, sport=sport, dport=dport,rport=rport)

def nats_from_xml(doc,groups):
    """
    Parse <dnats>, <snats>, <pnats> tags from NuFace XML
    """
    lsnat = []
    ldnat = []
    lpnat = []
    for snats in doc.getElementsByTagName('snats'):
       for snat in snats.getElementsByTagName('snat'):
            group = snat.getAttribute('group')
            if not group or groups[int(group)] == 1:
                id, name, srcnet, dstnet, proto, sport, dport, rewrite, rport, icmptype = _parse_snat_dnat(snat)
                snat = Snat(srcnet, dstnet, proto, rewrite, id = int(id), sport = sport, dport = dport)
                snat.icmptype = icmptype
                snat.set_name(name)
                lsnat.append(snat)

    for dnats in doc.getElementsByTagName('dnats'):
        for dnat in dnats.getElementsByTagName('dnat'):
            group = dnat.getAttribute('group')
            if not group or groups[int(group)] == 1:
                id, name, srcnet, dstnet, proto, sport, dport, rewrite, rport, icmp = _parse_snat_dnat(dnat)
                dnat = Dnat(srcnet, dstnet, proto, rewrite, id = int(id), sport = sport, dport = dport,rport=rport)
                dnat.icmptype = icmp
                dnat.set_name(name)
                ldnat.append(dnat)

    for pnats in doc.getElementsByTagName('pnats'):
        for pnat in pnats.getElementsByTagName('pnat'):
            group = pnat.getAttribute('group')
            if not group or groups[int(group)] == 1:
                id, name, srcnet, dstnet, proto, sport, dport, rewrite, rport, icmp = _parse_snat_dnat(pnat)
                p = Pnat(srcnet, dstnet, proto, dstnet, id = int(id), sport = sport, dport = dport, rport=rport)
                p.icmptype = icmp
                p.set_name(name)
                lpnat.append(p)
    return lsnat,ldnat,lpnat


def _parse_snat_dnat(node):
    """Parse an nat rule: snat, dnat or pnat"""
    name = node.getAttribute('name')
    srcnet = node.getAttribute('srcnet')
    dstnet = node.getAttribute('dstnet')
    sport = node.getAttribute('sport')
    dport = node.getAttribute('dport')
    rewrite = node.getAttribute('rewritetoaddr')
    icmptype = node.getAttribute('icmptype')
    rewriteport = node.getAttribute('rewritetoport')

    proto = node.getAttribute('proto')
    id = node.getAttribute('ID')
    try:
        proto = ip_protocols_rev[proto]
    except KeyError:
        proto = None
    return [id, name, srcnet, dstnet, proto, sport, dport, rewrite,
        rewriteport, icmptype]