Changeset 166 for etherws


Ignore:
Timestamp:
06/27/12 01:38:46 (12 years ago)
Author:
atzm
Message:
  • restructure
File:
1 edited

Legend:

Unmodified
Added
Removed
  • etherws/trunk/etherws.py

    r165 r166  
    5151import traceback 
    5252 
    53 import pytun 
    5453import websocket 
    5554import tornado.web 
     
    5857import tornado.httpserver 
    5958 
     59from pytun import TunTapDevice, IFF_TAP, IFF_NO_PI 
     60 
    6061 
    6162class DebugMixIn(object): 
    62     def dprintf(self, msg, func): 
     63    def dprintf(self, msg, func=lambda: ()): 
    6364        if self._debug: 
    6465            prefix = '[%s] %s - ' % (time.asctime(), self.__class__.__name__) 
     
    9394 
    9495 
    95 class SwitchingTable(DebugMixIn): 
     96class FDB(DebugMixIn): 
    9697    def __init__(self, ageout=300, debug=False): 
    9798        self._ageout = ageout 
    9899        self._debug = debug 
    99         self._table = {} 
    100  
    101     def learn(self, frame, port): 
    102         mac = frame.src_mac 
    103         vid = frame.vid 
    104  
    105         if vid not in self._table: 
    106             self._table[vid] = {} 
    107  
    108         self._table[vid][mac] = {'time': time.time(), 'port': port} 
    109         self.dprintf('learned: [%d] %s\n', 
    110                      lambda: (vid, mac.encode('hex'))) 
     100        self._dict = {} 
    111101 
    112102    def lookup(self, frame): 
     
    114104        vid = frame.vid 
    115105 
    116         group = self._table.get(vid, None) 
     106        group = self._dict.get(vid, None) 
    117107        if not group: 
    118108            return None 
     
    123113 
    124114        if time.time() - entry['time'] > self._ageout: 
    125             del self._table[vid][mac] 
    126             if not self._table[vid]: 
    127                 del self._table[vid] 
     115            del self._dict[vid][mac] 
     116            if not self._dict[vid]: 
     117                del self._dict[vid] 
    128118            self.dprintf('aged out: [%d] %s\n', 
    129119                         lambda: (vid, mac.encode('hex'))) 
     
    132122        return entry['port'] 
    133123 
     124    def learn(self, port, frame): 
     125        mac = frame.src_mac 
     126        vid = frame.vid 
     127 
     128        if vid not in self._dict: 
     129            self._dict[vid] = {} 
     130 
     131        self._dict[vid][mac] = {'time': time.time(), 'port': port} 
     132        self.dprintf('learned: [%d] %s\n', 
     133                     lambda: (vid, mac.encode('hex'))) 
     134 
    134135    def delete(self, port): 
    135         for vid in self._table.keys(): 
    136             for mac in self._table[vid].keys(): 
    137                 if self._table[vid][mac]['port'] is port: 
    138                     del self._table[vid][mac] 
     136        for vid in self._dict.keys(): 
     137            for mac in self._dict[vid].keys(): 
     138                if self._dict[vid][mac]['port'] is port: 
     139                    del self._dict[vid][mac] 
    139140                    self.dprintf('deleted: [%d] %s\n', 
    140141                                 lambda: (vid, mac.encode('hex'))) 
    141             if not self._table[vid]: 
    142                 del self._table[vid] 
    143  
    144  
    145 class TapHandler(DebugMixIn): 
    146     READ_SIZE = 65535 
    147  
    148     def __init__(self, dev, debug=False): 
     142            if not self._dict[vid]: 
     143                del self._dict[vid] 
     144 
     145 
     146class SwitchingHub(DebugMixIn): 
     147    def __init__(self, fdb, debug=False): 
     148        self._fdb = fdb 
    149149        self._debug = debug 
    150         self._clients = [] 
    151         self._table = SwitchingTable(debug=debug) 
    152         self._tap = pytun.TunTapDevice(dev, pytun.IFF_TAP | pytun.IFF_NO_PI) 
    153         self._tap.up() 
    154         self.register_client(self) 
    155  
    156     def fileno(self): 
    157         return self._tap.fileno() 
    158  
    159     def register_client(self, client): 
    160         self._clients.append(client) 
    161  
    162     def unregister_client(self, client): 
    163         self._table.delete(client) 
    164         self._clients.remove(client) 
    165  
    166     def write_message(self, message, binary=False): 
    167         self._tap.write(message) 
    168  
    169     def write(self, caller, message): 
    170         frame = EthernetFrame(message) 
    171  
    172         self._table.learn(frame, caller) 
    173  
    174         if not frame.multicast: 
    175             dst = self._table.lookup(frame) 
    176  
    177             if dst: 
    178                 dst.write_message(frame.data, True) 
    179                 self.dprintf('sent unicast: [%d] %s -> %s\n', 
    180                              lambda: (frame.vid, 
    181                                       frame.src_mac.encode('hex'), 
    182                                       frame.dst_mac.encode('hex'))) 
    183                 return 
    184  
    185         clients = self._clients[:] 
    186         clients.remove(caller) 
    187  
    188         for c in clients: 
    189             c.write_message(frame.data, True) 
    190  
     150        self._ports = [] 
     151 
     152    def register_port(self, port): 
     153        self._ports.append(port) 
     154 
     155    def unregister_port(self, port): 
     156        self._fdb.delete(port) 
     157        self._ports.remove(port) 
     158 
     159    def forward(self, src_port, frame): 
     160        try: 
     161            self._fdb.learn(src_port, frame) 
     162 
     163            if not frame.multicast: 
     164                dst_port = self._fdb.lookup(frame) 
     165 
     166                if dst_port: 
     167                    self._unicast(frame, dst_port) 
     168                    return 
     169 
     170            self._broadcast(frame, src_port) 
     171 
     172        except:  # ex. received invalid frame 
     173            traceback.print_exc() 
     174 
     175    def _unicast(self, frame, port): 
     176        port.write_message(frame.data, True) 
     177        self.dprintf('sent unicast: [%d] %s -> %s\n', 
     178                     lambda: (frame.vid, 
     179                              frame.src_mac.encode('hex'), 
     180                              frame.dst_mac.encode('hex'))) 
     181 
     182    def _broadcast(self, frame, *except_ports): 
     183        ports = self._ports[:] 
     184        for port in except_ports: 
     185            ports.remove(port) 
     186        for port in ports: 
     187            port.write_message(frame.data, True) 
    191188        self.dprintf('sent broadcast: [%d] %s -> %s\n', 
    192189                     lambda: (frame.vid, 
     
    194191                              frame.dst_mac.encode('hex'))) 
    195192 
     193 
     194class TapHandler(DebugMixIn): 
     195    READ_SIZE = 65535 
     196 
     197    def __init__(self, switch, dev, debug=False): 
     198        self._switch = switch 
     199        self._dev = dev 
     200        self._debug = debug 
     201        self._tap = None 
     202 
     203    @property 
     204    def closed(self): 
     205        return not self._tap 
     206 
     207    def open(self): 
     208        if not self.closed: 
     209            raise ValueError('already opened') 
     210        self._tap = TunTapDevice(self._dev, IFF_TAP | IFF_NO_PI) 
     211        self._tap.up() 
     212        self._switch.register_port(self) 
     213 
     214    def close(self): 
     215        if self.closed: 
     216            raise ValueError('I/O operation on closed tap') 
     217        self._switch.unregister_port(self) 
     218        self._tap.close() 
     219        self._tap = None 
     220 
     221    def fileno(self): 
     222        if self.closed: 
     223            raise ValueError('I/O operation on closed tap') 
     224        return self._tap.fileno() 
     225 
     226    def write_message(self, message, binary=False): 
     227        if self.closed: 
     228            raise ValueError('I/O operation on closed tap') 
     229        self._tap.write(message) 
     230 
    196231    def __call__(self, fd, events): 
     232        try: 
     233            self._switch.forward(self, EthernetFrame(self._read())) 
     234            return 
     235        except: 
     236            traceback.print_exc() 
     237        tornado.ioloop.IOLoop.instance().stop() 
     238 
     239    def _read(self): 
     240        if self.closed: 
     241            raise ValueError('I/O operation on closed tap') 
    197242        buf = [] 
    198  
    199243        while True: 
    200             data = self._tap.read(self.READ_SIZE) 
    201  
    202             if data: 
    203                 buf.append(data) 
    204  
    205             if len(data) < self.READ_SIZE: 
     244            buf.append(self._tap.read(self.READ_SIZE)) 
     245            if len(buf[-1]) < self.READ_SIZE: 
    206246                break 
    207  
    208         self.write(self, ''.join(buf)) 
     247        return ''.join(buf) 
    209248 
    210249 
    211250class EtherWebSocketHandler(tornado.websocket.WebSocketHandler, DebugMixIn): 
    212     def __init__(self, app, req, tap, debug=False): 
     251    def __init__(self, app, req, switch, debug=False): 
    213252        super(EtherWebSocketHandler, self).__init__(app, req) 
    214         self._tap = tap 
     253        self._switch = switch 
    215254        self._debug = debug 
    216255 
    217256    def open(self): 
    218         self._tap.register_client(self) 
     257        self._switch.register_port(self) 
    219258        self.dprintf('connected: %s\n', lambda: self.request.remote_ip) 
    220259 
    221260    def on_message(self, message): 
    222         self._tap.write(self, message) 
     261        self._switch.forward(self, EthernetFrame(message)) 
    223262 
    224263    def on_close(self): 
    225         self._tap.unregister_client(self) 
     264        self._switch.unregister_port(self) 
    226265        self.dprintf('disconnected: %s\n', lambda: self.request.remote_ip) 
    227266 
    228267 
    229268class EtherWebSocketClient(DebugMixIn): 
    230     def __init__(self, tap, url, user=None, passwd=None, debug=False): 
    231         self._sock = None 
    232         self._tap = tap 
     269    def __init__(self, switch, url, user=None, passwd=None, debug=False): 
     270        self._switch = switch 
    233271        self._url = url 
    234272        self._debug = debug 
     273        self._sock = None 
    235274        self._options = {} 
    236275 
     
    249288        self._sock = websocket.WebSocket() 
    250289        self._sock.connect(self._url, **self._options) 
     290        self._switch.register_port(self) 
    251291        self.dprintf('connected: %s\n', lambda: self._url) 
    252292 
     
    254294        if self.closed: 
    255295            raise websocket.WebSocketException('already closed') 
     296        self._switch.unregister_port(self) 
    256297        self._sock.close() 
    257298        self._sock = None 
     
    276317            data = self._sock.recv() 
    277318            if data is not None: 
    278                 self._tap.write(self, data) 
     319                self._switch.forward(self, EthernetFrame(data)) 
    279320                return 
    280321        except: 
     
    389430                              load_htpasswd(args.htpasswd)) 
    390431 
    391     tap = TapHandler(args.device, debug=args.debug) 
     432    fdb = FDB(debug=args.debug) 
     433    switch = SwitchingHub(fdb, debug=args.debug) 
     434    tap = TapHandler(switch, args.device, debug=args.debug) 
    392435    app = tornado.web.Application([ 
    393         (args.path, handler, {'tap': tap, 'debug': args.debug}), 
     436        (args.path, handler, {'switch': switch, 'debug': args.debug}), 
    394437    ]) 
    395438    server = tornado.httpserver.HTTPServer(app, ssl_options=ssl_options) 
     439 
     440    tap.open() 
    396441    server.listen(args.port, address=args.address) 
    397442 
     
    419464        args.passwd = getpass.getpass() 
    420465 
    421     tap = TapHandler(args.device, debug=args.debug) 
    422     client = EtherWebSocketClient(tap, args.uri, 
     466    fdb = FDB(debug=args.debug) 
     467    switch = SwitchingHub(fdb, debug=args.debug) 
     468    tap = TapHandler(switch, args.device, debug=args.debug) 
     469    client = EtherWebSocketClient(switch, args.uri, 
    423470                                  args.user, args.passwd, args.debug) 
    424471 
    425     tap.register_client(client) 
     472    tap.open() 
    426473    client.open() 
    427474 
Note: See TracChangeset for help on using the changeset viewer.