# Copyright(C) 2005 INL
# Written by Jean Gillaux <jean@inl.fr>
#
# 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/>.
# classes for iptables NAT rules

from nupyf.net_defs import ip_protocols, PROTO_TCP, PROTO_UDP, PROTO_ICMP
from nupyf.ipt_func import gen_rule_list
from IPy import IP

INTERNET_IPV4 = IP("0.0.0.0/0")

class GenNatIpt(object):
    """Produce iptables nat rules"""
    def __init__(self, rule):
        self._r = rule
        self._chain = ''
        self._target = ''
        self._targ = ''
        self._ifacearg = '-i'
        self.accept_vpn = True

    def _port_range(self, ports):
        if ports['l'] == ports['h'] and ports['h'] != -1:
            return [str(ports['h'])]
        elif ports['l'] != -1 and ports['h'] != -1:
            return [str(ports['l']),str(ports['h'])]

    def __str__(self):
        l = ["-A %s" % (self._chain)]
        proto = self._r._proto
        if self._r._iface and self._ifacearg:
            l.append('%s %s' % (self._ifacearg, self._r._iface.split(':')[0]))
        if self._r.proto_defined():
            s_proto = ip_protocols[proto]
            l.append("-p %s" % s_proto)
        if self._r._src and self._r._src != INTERNET_IPV4:
            l.append("-s %s" % (self._r._src))
        if self._r._dst and self._r._dst != INTERNET_IPV4:
            l.append("-d %s" % (self._r._dst))
        if proto in [PROTO_TCP, PROTO_UDP]:
            if self._r.port_defined(self._r._sport):
                l.append("--sport %s" % (':'.join(self._port_range(self._r._sport))))
            if self._r.port_defined(self._r._dport):
                l.append("--dport %s" % (':'.join(self._port_range(self._r._dport))))
        if proto == PROTO_ICMP:
            l.append("--icmp-type %s" % self._r.icmptype)
        l.append("-j %s" % self._target)
        to = self._r._to
        if self._r.port_defined(self._r._rport):
            to = "%s:%s" % (to,"-".join(self._port_range(self._r._rport)))
        if self._targ:
            l.append("%s %s" % (self._targ,to))
        if not self.accept_vpn:
            l.append("-m policy --pol none --dir out")
        return gen_rule_list(l)

class SnatIpt(GenNatIpt):
    def __init__(self, rule):
        super(SnatIpt, self).__init__(rule)
        self._chain = 'POSTROUTING'
        self._target = 'SNAT'
        self._targ = "--to-source"
        self._ifacearg = '-o'
        self.accept_vpn = False

class DnatIpt(GenNatIpt):
    def __init__(self, rule):
        super(DnatIpt, self).__init__(rule)
        self._chain = 'PREROUTING'
        self._target = 'DNAT'
        self._targ = "--to-destination"

class PnatIpt(GenNatIpt):
    def __init__(self, rule):
        super(PnatIpt, self).__init__(rule)
        self._chain = 'PREROUTING'
        self._target = 'DNAT'
        self._targ = '--to-destination'