############################################################################
##
## Copyright (c) 2000-2001 BalaBit IT Ltd, Budapest, Hungary
## All rights reserved.
##
## $Id: NAT.py,v 1.3.2.2 2001/09/04 09:42:56 bazsi Exp $
##
############################################################################

"""Module defining interfaces for network address translation.

This module defines 'AbstractNAT' an abstract interface for performing
address translation for client addresses.
"""

from Zorp import *
from SockAddr import SockAddrInet
from Domain import InetDomain
from whrandom import choice

class AbstractNAT:
	"""Abstract class encapsulating interface for address translation.
	
	This class encapsulates an interface for application level network
	address translation. This NAT is not the same you are used to with
	regard to packet filters, it basically specifies/modifies outgoing
	source/destination addresses just before Zorp connects to the
	server.
	
	Source and destination NATs can be specified when a 'Service' is
	created.
	
	Currently it is used by the 'ConnectChainer' class just before 
	connecting to the server.
	"""
	
	def performTranslation(self, session, addr):
		"""Function to actually perform address translation.
		
		This function is called when a session is to be connected
		to its server endpoint. It should return the address to
		bind to before connecting.
		
		Arguments
		
		  self -- this instance
		  
		  session -- session which is about to connect
		  
		  addr -- address to translate
		
		Returns
		
		  A SockAddr instance to bind to.
		"""
		raise NotImplementedError

class ForgeClientSourceNAT(AbstractNAT):
	"""Class to fake original client address for outgoing connections.
	
	This class implements the 'AbstractNAT' interface to perform address
	translation. It effectively uses the client address, so the connection
	seem to come from the original client instead of the firewall.
	
	It can be used when the original client address has some value (for
	example WWW and other servers in the DMZ performing address based
	access control).
	
	"""
	def performTranslation(self, session, addr):
		"""Function to actually perform address translation.
		
		This function returns the original client address with
		a dynamically allocated port number.
		
		Arguments
		
		  self -- this instance
		  
		  session -- session which is about to connect
		  
		  addr -- address to translate
		  
		Returns
		
		  The address to bind to
		"""
		
		return SockAddrInet(session.client_address.ip_s, 0)

class OneToOneNAT(AbstractNAT):
	"""Class to translate from a given address range to another.
	
	This class performs 1:1 address translation from the given source
	domain to the also given destination domain. If the source address
	is outside of the given source address range, DACException is raised.
	The supplied address ranges must have the same size.
	
	Instances of this class can be used to direct traffic destined to
	a block of IPs to another block, like having a lot of Web servers
	in your DMZ with dedicated IP aliases on your firewall.
	
	Attributes
	
	  from_domain  -- source subnet (InetDomain instance)
	  
	  to_domain    -- destination subnet (InetDomain instance)

	  default_reject -- whether to reject all connections not within the required range
	"""
        def __init__(self, from_domain, to_domain, default_reject=TRUE):
		"""Constructor to initialize a OneToOneNAT instance
		
		This constructor initializes a OneToOneNAT instance by
		storing arguments in attributes. Arguments must be
		InetDomain instances specifying two non-overlapping IP subnets
		with the same size.
		
		Arguments
		
		  from_domain  -- source address range to translate addresses from
		  
		  to_domain    -- destination address range to translate addresses into
		"""
		AbstractNAT.__init__(self)
		self.from_domain = from_domain
		self.to_domain = to_domain
		self.default_reject = default_reject
		if from_domain.mask_bits != to_domain.mask_bits:
			raise ValueError, 'OneToOneNAT requires two domains of the same size'
		
	def performTranslation(self, session, addr):
		"""Function to actually perform 1:1 NAT
		
		This function actually performs the required changes to
		project an address in the source address range to the
		destination address range.
		
		Arguments
		
		  self -- this instance
		  
		  session -- session we belong to
		  
		  addr -- address to translate
		  
		Returns
		
		  This function returns the calculated new address.
		"""
		if addr < self.from_domain:
			ip = (addr.ip & ~self.to_domain.mask) + (to_domain.ip & self.to_domain.mask)
			return SockAddrInet(inet_ntoa(ip), 0)
		if default_reject:
			return DACException, 'IP not within the required range.'
		else:
			return addr
		
class RandomNAT(AbstractNAT):
	"""Class to choose one of several addresses randomly.
	
	This class chooses one of the supplied addresses randomly.
	It can be used for load-balancing several lines by binding
	to different interfaces for each session.
	
	Attributes
	
	  addresses -- sequence of addresses to choose one from
	"""
	def __init__(self, addresses):
		"""Constructor to initialize a RandomNAT instance.
		
		This constructor initializes a RandomNAT instance by
		calling the inherited constructor and storing arguments
		in appropriate attributes.
		
		Arguments
		
		  addresses -- sequence of addresses to choose one from, 
		               each item in this sequence must be a SockAddr
			       (or derived class) instance.
		"""
		AbstractNAT.__init__(self)
		self.addresses = addresses
	
	def performTranslation(self, session, addr):
		"""Function to actually perform random-address translation.
		
		This function actually performs address translation by
		choosing one of the supplied addresses randomly.
		
		Arguments
		
		  self -- this instance
		  
		  session -- session we belong to
		  
		  addr -- address to translate
		  
		Returns
		
		  The translated address is returned.
		"""
		return choice(self.addresses)
