#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
SRXをエミュレーションするモジュール
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Juniper SRXのpolicy追加をテストするための
モジュールです。
:copyright:
:license:
"""

import re
from collections.abc import Sequence
from ipaddress import IPv4Network, IPv4Address


SourceAddress = "source_address"
DestinationAddress = "destination_address"
Application = "application"
ANY = "any"

Permit = "permit"
Deny = "deny"
Reject = "reject"
Log = "log"
Count = "count"

class PsuedoSRX(object):
    """PsuedoSRX

    SRXのインスタンスをエミュレートします
    """
    def __init__(self, device_name):
        super(PsuedoSRX, self).__init__()
        self.device_name = device_name
        self.address_books = {}
        self.set_address_book(ANY)
        self.applications = {}
        self.set_applications(ANY)
        self.zones = {}
        self.policies = []

    def set_address_book(self, book_name, nw_address=None):
        book = self.__get_book(book_name)
        if nw_address:
            book.append(nw_address)

    def set_zone(self, zone_name):
        zone = Zone(zone_name)
        self.zones[zone.name] = zone

    def set_applications(self, app_name, protocol=None,
                         src_or_dst="destination-port", port=None):
        app = self.__get_app(app_name)
        if protocol:
            app.protocol = protocol
        if port:
            app.port = port
        if src_or_dst == "source-port":
            app.is_src = True

    def set_policy_with_addressbook(self, policy_name, src_or_dst, book_name):
        policy = self.__get_policy(policy_name)
        book = self.__get_book(book_name)
        policy.add_address_book(src_or_dst, book)

    def set_policy_with_application(self, policy_name, app_name):
        policy = self.__get_policy(policy_name)
        app = self.__get_app(app_name)
        policy.add_application(app)

    def set_policy_with_action(self, policy_name, action):
        policy = self.__get_policy(policy_name)
        policy.add_action(action)

    def insert_policy(self, policy_name, insert_name, before_or_after="before"):
        policy = self.__get_policy(policy_name)
        self.policies.remove(policy)
        if before_or_after == "before":
            target = self.__get_policy(insert_name)
            i = self.policies.index(target)
        else:
            target = self.__get_policy(insert_name)
            i = self.policies.index(target) + 1
        self.policies.insert(i, policy)

    def packet_in(self, protocol, src_addr, src_port, dest_addr, dest_port):
        # Packetが入ってきたときの処理を記載する
        for policy in self.policies:
            if policy.lookup(protocol, src_addr, src_port, dest_addr, dest_port):
                return policy.action()
        return [Deny]

    def __get_policy(self, policy_name):
        policy = None
        for p in self.policies:
            if p.policy_name == policy_name:
                policy = p
                break
        if not policy:
            policy = Policy(policy_name)
            self.policies.append(policy)
        return policy

    def __get_book(self, book_name):
        if book_name in self.address_books:
            return self.address_books[book_name]
        book = AddressBook(book_name)
        self.address_books[book_name] = book
        return book

    def __get_app(self, app_name):
        if app_name in self.applications:
            return self.applications[app_name]
        app = Application(app_name)
        self.applications[app_name] = app
        return app

    def __repr__(self):
        """

        jsonを返すようにしたい
        """
        return "{0}:{1}".format(self.device_name, self.applications)

class Zone(object):
    """Zone

    インターフェースをエミュレートするのが困難なので
    Zoneも実体としては定義していません。

    """
    def __init__(self, name):
        super(Zone, self).__init__()
        self.name = name
        self.address_books = []

    def attach(self, address_book):
        self.address_books.append(address_book)


class AddressBook(Sequence):
    """AddressBook

    junos SRX のAddressBookです。
    :param book_name: this address-book's name
    :param nw: the nw that this address-book has.
    """

    def __init__(self, book_name, nw=None):
        super(AddressBook, self).__init__()
        self.book_name = book_name
        if book_name == ANY:
            nw = '0.0.0.0/0'
        self.nws = []
        if nw:
            self.nws.append(IPv4Network(nw))

    def append(self, nw):
        self.nws.append(IPv4Network(nw))

    def __getitem__(self, key):
        return self.nws[key]

    def __len__(self):
        return len(self.nws)

    def __contains__(self, item):
        n = IPv4Address(item)
        for nw in self.nws:
            if n in nw:
                return True
        return False

    def __repr__(self):
        s = 'AddressBook {0}: '.format(self.book_name)
        for _nw in self.nws:
            s += '{0} '.format(_nw.exploded)
        return s


class Application(object):
    """Application"""
    port_pattern_range = re.compile('(\d+)-(\d+)')
    port_pattern_simple = re.compile('\d+')
    port_pattern_junos_name = re.compile('junos-[a-zA-Z]+')
    junos_app_name_mapper = {'junos-ssh':22, 'junos-http':80}

    @classmethod
    def __to_number(self, junos_app_name):
        return self.junos_app_name_mapper[junos_app_name]

    def __init__(self, app_name):
        super(Application, self).__init__()
        self.app_name = app_name
        self.is_any = False
        self.is_src = False
        if app_name == ANY:
            self.is_any = True
        self._protocol = ""
        self._port = None

    @property
    def protocol(self):
        return self._protocol

    @protocol.setter
    def protocol(self, val):
        self._protocol = val

    @property
    def port(self):
        return self._port

    @port.setter
    def port(self, ports):
        m = self.port_pattern_range.match(ports)
        if m:
            self._port = range(int(m.group(1)), int(m.group(2))+1)
            return
        m = self.port_pattern_simple.match(ports)
        if m:
            self._port = []
            self._port.append(int(m.group(0)))
            return
        if self.port_pattern_junos_name.match(ports):
            self._port = []
            self._port.append(self.__to_number(ports))
            return
        else:
            raise

    def match(self, proto, src_port=None, dest_port=None):
        if self.is_any:
            return True
        if (proto == self._protocol) and self._port:
            if dest_port in self._port:
                return True
            if src_port in self._port:
                return True
        return False

    def __repr__(self):
        return "{0} {1} {2}".format(self.app_name, self._protocol, self._port)


class Policy(object):
    """Policy
    """
    def __init__(self, policy_name):
        super(Policy, self).__init__()
        self.policy_name = policy_name
        self.source_addresses = []
        self.destination_addresses = []
        self.applications = []
        self.actions = []

    def add_from_zone(self, from_zone):
        self.from_zone = from_zone

    def add_to_zone(self, to_zone):
        self.to_zone = to_zone

    def add_address_book(self, src_or_dst, book):
        if src_or_dst == "source-address":
            self.source_addresses.append(book)
        else:
            self.destination_addresses.append(book)

    def add_application(self, app):
        self.applications.append(app)

    def add_action(self, action):
        self.actions.append(action)

    def lookup(self, protocol, src_addr, src_port, dest_addr, dest_port):
        if not self.__is_in_source_addresses(src_addr):
            return False
        if not self.__is_in_destination_addresses(dest_addr):
            return False
        if not self.__is_in_application(protocol, src_port, dest_port):
            return False
        return True

    def action(self):
        return self.actions

    def __is_in_source_addresses(self, src_addr):
        for book in self.source_addresses:
            if src_addr in book:
                return True
        return False

    def __is_in_destination_addresses(self, dest_addr):
        for book in self.destination_addresses:
            if dest_addr in book:
                return True
        return False

    def __is_in_application(self, protocol, src_port, dest_port):
        for app in self.applications:
            if app.match(protocol, src_port, dest_port):
                return True
        return False

    def __repr__(self):
        return self.policy_name
