#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#                         Ethernet over IRC client
#
# depends on:
#   - python-2.7.2
#   - python-pytun-0.2
#   - tornado-2.3
#
# ===========================================================================
# Copyright (c) 2012, Atzm WATANABE <atzm@atzm.org>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# ===========================================================================
#
# $Id$

import os
import sys
import time
import errno
import socket
import argparse
import traceback

import tornado.ioloop
from pytun import TunTapDevice, IFF_TAP, IFF_NO_PI


class DebugMixIn(object):
    def dprintf(self, msg, func=lambda: ()):
        if self._debug:
            prefix = '[%s] %s - ' % (time.asctime(), self.__class__.__name__)
            sys.stderr.write(prefix + (msg % func()))


class EthernetFrame(object):
    def __init__(self, data):
        self.data = data

    @property
    def dst_multicast(self):
        return ord(self.data[0]) & 1

    @property
    def src_multicast(self):
        return ord(self.data[6]) & 1

    @property
    def dst_mac(self):
        return self.data[:6]

    @property
    def src_mac(self):
        return self.data[6:12]

    @property
    def tagged(self):
        return ord(self.data[12]) == 0x81 and ord(self.data[13]) == 0

    @property
    def vid(self):
        if self.tagged:
            return ((ord(self.data[14]) << 8) | ord(self.data[15])) & 0x0fff
        return 0

    @staticmethod
    def format_mac(mac, sep=':'):
        return sep.join(b.encode('hex') for b in mac)


class FDB(DebugMixIn):
    class Entry(object):
        def __init__(self, port, ageout):
            self.port = port
            self._time = time.time()
            self._ageout = ageout

        @property
        def age(self):
            return time.time() - self._time

        @property
        def agedout(self):
            return self.age > self._ageout

    def __init__(self, ageout, debug=False):
        self._ageout = ageout
        self._debug = debug
        self._table = {}

    def _set_entry(self, vid, mac, port):
        if vid not in self._table:
            self._table[vid] = {}
        self._table[vid][mac] = self.Entry(port, self._ageout)

    def _del_entry(self, vid, mac):
        if vid in self._table:
            if mac in self._table[vid]:
                del self._table[vid][mac]
            if not self._table[vid]:
                del self._table[vid]

    def _get_entry(self, vid, mac):
        try:
            entry = self._table[vid][mac]
        except KeyError:
            return None

        if not entry.agedout:
            return entry

        self._del_entry(vid, mac)
        self.dprintf('aged out: port:%d; vid:%d; mac:%s\n',
                     lambda: (entry.port.number, vid, mac.encode('hex')))

    def each(self):
        for vid in sorted(self._table.iterkeys()):
            for mac in sorted(self._table[vid].iterkeys()):
                entry = self._get_entry(vid, mac)
                if entry:
                    yield (vid, mac, entry)

    def lookup(self, frame):
        mac = frame.dst_mac
        vid = frame.vid
        entry = self._get_entry(vid, mac)
        return getattr(entry, 'port', None)

    def learn(self, port, frame):
        mac = frame.src_mac
        vid = frame.vid
        self._set_entry(vid, mac, port)
        self.dprintf('learned: port:%d; vid:%d; mac:%s\n',
                     lambda: (port.number, vid, mac.encode('hex')))

    def delete(self, port):
        for vid, mac, entry in self.each():
            if entry.port.number == port.number:
                self._del_entry(vid, mac)
                self.dprintf('deleted: port:%d; vid:%d; mac:%s\n',
                             lambda: (port.number, vid, mac.encode('hex')))


class SwitchingHub(DebugMixIn):
    class Port(object):
        def __init__(self, number, interface):
            self.number = number
            self.interface = interface
            self.tx = 0
            self.rx = 0
            self.shut = False

        @staticmethod
        def cmp_by_number(x, y):
            return cmp(x.number, y.number)

    def __init__(self, fdb, debug=False):
        self.fdb = fdb
        self._debug = debug
        self._table = {}
        self._next = 1

    @property
    def portlist(self):
        return sorted(self._table.itervalues(), cmp=self.Port.cmp_by_number)

    def get_port(self, portnum):
        return self._table[portnum]

    def register_port(self, interface):
        try:
            self._set_privattr('portnum', interface, self._next)  # XXX
            self._table[self._next] = self.Port(self._next, interface)
            return self._next
        finally:
            self._next += 1

    def unregister_port(self, interface):
        portnum = self._get_privattr('portnum', interface)
        self._del_privattr('portnum', interface)
        self.fdb.delete(self._table[portnum])
        del self._table[portnum]

    def send(self, dst_interfaces, frame):
        portnums = (self._get_privattr('portnum', i) for i in dst_interfaces)
        ports = (self._table[n] for n in portnums)
        ports = (p for p in ports if not p.shut)
        ports = sorted(ports, cmp=self.Port.cmp_by_number)

        for p in ports:
            p.interface.write_message(frame.data)
            p.tx += 1

        if ports:
            self.dprintf('sent: port:%s; vid:%d; %s -> %s\n',
                         lambda: (','.join(str(p.number) for p in ports),
                                  frame.vid,
                                  frame.src_mac.encode('hex'),
                                  frame.dst_mac.encode('hex')))

    def receive(self, src_interface, frame):
        port = self._table[self._get_privattr('portnum', src_interface)]

        if not port.shut:
            port.rx += 1
            self._forward(port, frame)

    def _forward(self, src_port, frame):
        try:
            if not frame.src_multicast:
                self.fdb.learn(src_port, frame)

            if not frame.dst_multicast:
                dst_port = self.fdb.lookup(frame)

                if dst_port:
                    self.send([dst_port.interface], frame)
                    return

            ports = set(self.portlist) - set([src_port])
            self.send((p.interface for p in ports), frame)

        except:  # ex. received invalid frame
            traceback.print_exc()

    def _privattr(self, name):
        return '_%s_%s_%s' % (self.__class__.__name__, id(self), name)

    def _set_privattr(self, name, obj, value):
        return setattr(obj, self._privattr(name), value)

    def _get_privattr(self, name, obj, defaults=None):
        return getattr(obj, self._privattr(name), defaults)

    def _del_privattr(self, name, obj):
        return delattr(obj, self._privattr(name))


class TapHandler(DebugMixIn):
    READ_SIZE = 65535

    def __init__(self, ioloop, switch, dev, debug=False):
        self._ioloop = ioloop
        self._switch = switch
        self._dev = dev
        self._debug = debug
        self._tap = None

    @property
    def target(self):
        if self.closed:
            return self._dev
        return self._tap.name

    @property
    def closed(self):
        return not self._tap

    @property
    def address(self):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        try:
            return self._tap.addr
        except:
            return ''

    @property
    def netmask(self):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        try:
            return self._tap.netmask
        except:
            return ''

    @property
    def mtu(self):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        return self._tap.mtu

    @address.setter
    def address(self, address):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        self._tap.addr = address

    @netmask.setter
    def netmask(self, netmask):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        self._tap.netmask = netmask

    @mtu.setter
    def mtu(self, mtu):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        self._tap.mtu = mtu

    def open(self):
        if not self.closed:
            raise ValueError('Already opened')
        self._tap = TunTapDevice(self._dev, IFF_TAP | IFF_NO_PI)
        self._tap.up()
        self._ioloop.add_handler(self.fileno(), self, self._ioloop.READ)
        return self._switch.register_port(self)

    def close(self):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        self._switch.unregister_port(self)
        self._ioloop.remove_handler(self.fileno())
        self._tap.close()
        self._tap = None

    def fileno(self):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        return self._tap.fileno()

    def write_message(self, message, binary=False):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        self._tap.write(message)

    def __call__(self, fd, events):
        try:
            self._switch.receive(self, EthernetFrame(self._read()))
            return
        except:
            traceback.print_exc()
        self.close()

    def _read(self):
        if self.closed:
            raise ValueError('I/O operation on closed tap')
        buf = []
        while True:
            buf.append(self._tap.read(self.READ_SIZE))
            if len(buf[-1]) < self.READ_SIZE:
                break
        return ''.join(buf)


class EtherIRCHandler(DebugMixIn):
    READ_SIZE = 65535

    def __init__(self, ioloop, switch, server, nick, channel, debug=False):
        self._ioloop = ioloop
        self._switch = switch
        self._server = server
        self._nick = nick
        self._channel = channel
        self._debug = debug
        self._sock = None
        self._buffer = []

    @property
    def closed(self):
        return not self._sock

    def open(self):
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._sock.connect(self._server)
        self._sendraw('NICK %s' % self._nick)
        self._sendraw('USER %s 0 * :%s' % (self._nick, self._nick))
        self._sendraw('JOIN %s' % self._channel)

        self._sock.setblocking(False)
        self._ioloop.add_handler(self.fileno(), self, self._ioloop.READ)
        return self._switch.register_port(self)

    def close(self):
        self._switch.unregister_port(self)
        self._ioloop.remove_handler(self.fileno())
        self._sock.close()
        self._sock = None

    def fileno(self):
        return self._sock.fileno()

    def write_message(self, message):
        data = message.encode('hex')
        self._sendraw('PRIVMSG %s %s' % (self._channel, data))

    def _sendraw(self, data):
        self._sock.send(data + '\r\n')

    def __call__(self, fd, events):
        close = False

        try:
            while True:
                data = self._sock.recv(self.READ_SIZE)
                if not data:
                    close = True
                    break
                self._buffer.append(data)
        except socket.error as e:
            if e.errno != errno.EAGAIN:
                raise e

        lines = ''.join(self._buffer).split('\r\n')
        rest = lines.pop(-1)

        self._buffer = []
        if rest:
            self._buffer.append(rest)

        for line in lines:
            line = line.strip()
            prefix = None
            if line.startswith(':'):
                prefix, cmd, params = line.split(' ', 2)
            else:
                cmd, params = line.split(' ', 1)
            method = getattr(self, '_handle_%s' % cmd, self.__handle_default)
            method(prefix, cmd, params)

        if close:
            self.close()
            self.dprintf('connection closed\n')
            self._ioloop.stop()  # XXX: shutdown process when connection closed

    def __handle_default(self, prefix, cmd, params):
        self.dprintf('UNKNOWN %s %s %s\n', lambda: (prefix, cmd, params))

    def _handle_PING(self, prefix, cmd, params):
        self.dprintf('%s %s %s\n', lambda: (prefix, cmd, params))
        self._sendraw('PONG 0')

    def _handle_PRIVMSG(self, prefix, cmd, params):
        self.dprintf('%s %s %s\n', lambda: (prefix, cmd, params))
        to, data = params.split(' ', 1)
        try:
            message = data[1:].decode('hex')
            self._switch.receive(self, EthernetFrame(message))
        except:
            traceback.print_exc()


def _main():
    def daemonize(nochdir=False, noclose=False):
        if os.fork() > 0:
            sys.exit(0)

        os.setsid()

        if os.fork() > 0:
            sys.exit(0)

        if not nochdir:
            os.chdir('/')

        if not noclose:
            os.umask(0)
            sys.stdin.close()
            sys.stdout.close()
            sys.stderr.close()
            os.close(0)
            os.close(1)
            os.close(2)
            sys.stdin = open(os.devnull)
            sys.stdout = open(os.devnull, 'a')
            sys.stderr = open(os.devnull, 'a')

    parser = argparse.ArgumentParser()
    parser.add_argument('--device', action='append', default=[])
    parser.add_argument('--ageout', action='store', type=int, default=300)
    parser.add_argument('--foreground', action='store_true', default=False)
    parser.add_argument('--debug', action='store_true', default=False)

    parser.add_argument('host')
    parser.add_argument('port', type=int, default=6667)
    parser.add_argument('nick')
    parser.add_argument('channel')

    args = parser.parse_args()

    if args.ageout <= 0:
        raise ValueError('invalid ageout: %s' % args.ageout)

    if not (0 < args.port < 65535):
        raise ValueError('invalid port: %s' % args.port)

    ioloop = tornado.ioloop.IOLoop.instance()
    fdb = FDB(ageout=args.ageout, debug=args.debug)
    switch = SwitchingHub(fdb, debug=args.debug)

    client = EtherIRCHandler(ioloop, switch, (args.host, args.port),
                             args.nick, args.channel, args.debug)
    client.open()

    for dev in args.device:
        tap = TapHandler(ioloop, switch, dev, debug=args.debug)
        tap.open()

    if not args.foreground:
        daemonize()

    ioloop.start()


if __name__ == '__main__':
    _main()