/
econat.py
executable file
·648 lines (559 loc) · 21.3 KB
/
econat.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
#!/usr/bin/env python3
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
import asyncio
import socket
from collections import namedtuple
import time
import logging
import logging.handlers
import json
import daemon
import daemon.pidfile
import lockfile
import signal
import sys
import setproctitle
import aiopg
import psycopg2
version = "1.0-0"
servicename = "econat"
servicedesc = "econat activator"
class User(namedtuple('User', ('id', 'ip', 'speed', 'service', 'port', 'status'))):
def __new__(cls, id=None, ip=None, speed=None, service=None, port=None, status=None):
return super(User, cls).__new__(cls, id, ip, speed, service, port, status)
def dict_compare(d1, d2):
d1_keys = set(d1.keys())
d2_keys = set(d2.keys())
intersect_keys = d1_keys.intersection(d2_keys)
added = d1_keys - d2_keys
removed = d2_keys - d1_keys
modified = {o for o in intersect_keys if d1[o] != d2[o]}
return added, removed, modified
class DbInfo:
"""
radius=> select * from active_users;
id | ip | speed | service | port | status
----+------------+----------+---------+------------------------------------+--------
1 | 100.64.1.3 | 20971520 | 20Mbit | eltex-1-1 eth 100/2:100 | 1
2 | 100.64.3.2 | 10485760 | 10Mbit | eltex-1-1 eth 100/1:100 | 1
2 | 100.64.1.4 | 10485760 | 10Mbit | eltex-1-1 eth 100/1:100 | 1
3 | 100.64.2.2 | 20971520 | 20Mbit | access-1 GigabitEthernet0/0/21:100 | 0
"""
def __init__(self, loop, user='radius', password='radius', dbname='radius'):
self.logid = 'DbInfo'
self.log = logging.getLogger(self.logid)
self.loop = loop
self.user = user
self.password = password
self.dbname = dbname
self.queuetask = None
@property
def queue(self):
if not self.__queue:
self.queue = asyncio.Queue()
return self.__queue
@queue.setter
def queue(self, queue):
self.__queue = queue
self.queuetask = self.loop.create_task(self.listen())
def connect(self):
self.log.info('connecting to db %s', self.dbname)
return aiopg.connect(dsn='dbname={dbname} user={user} password={password}'.format(dbname=self.dbname, user=self.user, password=self.password))
def stop(self):
if self.queuetask:
self.queuetask.cancel()
self.loop.run_until_complete(asyncio.wait([self.queuetask]))
async def execute(self, query, args=None):
try:
async with self.connect() as db:
async with db.cursor() as cur:
await cur.execute(query, args)
return await cur.fetchall()
except psycopg2.Error as e:
self.log.error('execute: %s %s', query, e)
async def listen(self):
async def wd(db):
self.log.info("start econat_notify watchdog")
while True:
if db.closed:
self.log.info("wd: db colsed %s", db.closed)
raise Exception("db reconnect")
await asyncio.sleep(1)
async def l(db):
async with db.cursor() as cur:
await cur.execute("LISTEN econat_notify")
while True:
try:
msg = await db.notifies.get()
self.log.log(1, '%s', msg)
command = json.loads(msg.payload)
await self.queue.put(command)
except json.decoder.JSONDecodeError as e:
self.log.error('l: receive %s', e)
while True:
_wd = None
_l = None
try:
async with self.connect() as db:
_wd = asyncio.ensure_future(wd(db))
_l = asyncio.ensure_future(l(db))
_wd.add_done_callback(lambda fut: _l.cancel())
self.log.info('listen econat_notify')
await asyncio.gather(_wd, _l)
except psycopg2.Error as e:
self.log.error('listen: %s', e)
await asyncio.sleep(1)
except asyncio.CancelledError:
self.log.info('close econat_notify')
break
except Exception:
self.log.info('listen: db closed. Reconnecting')
await asyncio.sleep(1)
async def rid(self):
data = await self.execute('select user_id, service_id from user_service where status = 1')
if data:
return dict(data)
async def userlist(self, user_tmpl, identificator, users):
if users:
fetchall = await self.execute("select %s from active_users where %s in ('%s')"
% (','.join(user_tmpl._fields), identificator, "','".join(map(str, users)))
)
self.log.log(1, "ret %s", fetchall)
return map(user_tmpl._make, fetchall)
return []
class TestUsers:
def stop(self):
pass
def rid(self):
userlist = {}
for i in range(1, 16):
userlist[i] = (i % 4)+1
return userlist
def userlist(self, user_tmpl, identificator, users):
userlist = []
for i in users:
userlist.append(User(id=i,ip='100.%d.%d.%d' % (64 + i % (65536 * 256) // 65536, i % 65536 // 256, i % 256), speed='%dM' % ((i % 4) + 1), service=((i % 4) + 1), port='-'))
return userlist
class EcoNat:
user = User
def __init__(self, loop, server = '192.168.100.200', port = 2225):
self.logid = 'EcoNat'
self.log = logging.getLogger(self.logid)
self.loop = loop
self.server = server
self.port = port
def stop(self):
pass
async def rid(self):
data = await self.ask(b'testRID\n')
return dict(map(int, i.split(b'-')) for i in data.split())
async def processusers(self, action_users={}):
mlist = []
mlist.append(self.addformat(action_users.get('add')))
mlist.append(self.delformat(action_users.get('del')))
data = ''.join(mlist).encode()
if data:
await self.ask(data, read=False)
async def ask(self, message, read=True):
data = b''
try:
self.log.debug('open connection to %s:%d', self.server, self.port)
reader, writer = await asyncio.open_connection(self.server,self.port, loop=self.loop)
self.log.debug('send %s', message)
writer.write(message)
if read:
data = await reader.readline()
self.log.debug('recive %s', data)
self.log.debug('close connection to %s:%d', self.server, self.port)
writer.close()
except OSError as e:
self.log.error('testRID connecting to %s:%d %s', self.server, self.port, e)
return data
def addformat(self, userlist):
return ''.join('add\t{0}\t{{oid}} LIM{1}/LIM{1} {2}, // RULE{3}\n'.format(user.id, user.speed, user.ip, user.service) for user in userlist)
def delformat(self, useridlist):
return ''.join('remove\t{0}\t\n'.format(user.id) for user in useridlist)
from pyrad import dictionary, packet
class RadiusHandler:
def __init__(self, loop, message):
self.logid = 'RadiusHandler{}'.format(message.get(1))
self.log = logging.getLogger(self.logid)
self.loop = loop
self.message = message
self.transport = None
self.data = None
self.scheduler = None
self.retransmit = 3
def send(self):
if self.retransmit:
self.log.debug('send %s', self.message)
self.retransmit -= 1
if not self.data:
self.data = self.message.RequestPacket()
self.transport.sendto(self.data)
self.scheduler = self.loop.call_later(5, self.send)
else:
self.transport.close()
def connection_made(self, transport):
self.log.debug('Connection made')
self.transport = transport
self.send()
def datagram_received(self, data, addr):
pkt = self.message.CreateReply(packet=data)
self.log.info("Received: %d", pkt.code)
if self.scheduler:
self.scheduler.cancel()
self.transport.close()
def error_received(self, exc):
self.log.error('Error received: %s', exc)
def connection_lost(self, exc):
self.log.debug("Socket closed")
class RadiusClient:
raddict = dictionary.Dictionary("/opt/econat2db/dictionary.rfc2865")
user = namedtuple('User', ('ip'))
def __init__(self, loop, server='192.168.100.200', port=1812, secret=b''):
self.loop = loop
self.logid = 'RadiusClient'
self.log = logging.getLogger(self.logid)
self.server = server
self.port = port
self.secret = secret
def stop(self):
pass
async def processusers(self, action_users={}):
for action, userlist in action_users.items():
for user in userlist:
if user.ip:
self.log.info('send radius_disconnect for %s', user.ip)
message = packet.AcctPacket(code=packet.DisconnectRequest, dict=self.raddict, secret=self.secret)
message['User-Name'] = user.ip
await self.loop.create_datagram_endpoint(lambda: RadiusHandler(self.loop, message), remote_addr=(self.server, self.port))
import os
class UnixSocket:
def __init__(self, loop, file=None):
self.logid = 'UnixSocket'
self.log = logging.getLogger(self.logid)
self.filename = file
self.socket = socket.socket( socket.AF_UNIX, socket.SOCK_DGRAM )
if os.path.exists(self.filename):
os.unlink(self.filename)
self.socket.bind(self.filename)
os.chmod(self.filename,0o1777)
self.loop = loop
self.transport = self.loop.create_datagram_endpoint(lambda: self, sock=self.socket)
self.loop.run_until_complete(self.transport)
@property
def queue(self):
if not self.__queue:
self.queue = asyncio.Queue()
return self.__queue
@queue.setter
def queue(self, queue):
self.__queue = queue
def connection_made(self, transport):
pass
def datagram_received(self, data, addr):
self.log.log(1, 'notify %s', data)
try:
command = json.loads(data.decode())
self.queue.put_nowait(command)
except json.decoder.JSONDecodeError as e:
self.log.error('receive %s %s', e, data)
def stop(self):
if os.path.exists(self.filename):
os.unlink(self.filename)
if self.transport:
self.transport.close()
#import stat
#import fcntl
#import pyinotify
#class FifoPipe(pyinotify.ProcessEvent):
# def my_init(self, loop, file=None):
# self.loop = loop
# if not file:
# raise ValueError("file keyword argument must be provided")
#
# self.filename = file
#
# if not os.path.exists(self.filename):
# os.mkfifo(self.filename)
#
# if not stat.S_ISFIFO(os.stat(self.filename).st_mode):
# raise TypeError("File %s is not a fifo file" % self.filename)
#
# self.fd = open(os.open(self.filename, os.O_RDONLY|os.O_NONBLOCK), 'r')
# self.queue = None
# wm = pyinotify.WatchManager()
# wm.add_watch(self.filename, pyinotify.IN_MODIFY)
# notifier = pyinotify.AsyncioNotifier(wm, self.loop, default_proc_fun=self)
#
# def process_IN_MODIFY(self, event):
# for data in self.fd:
# self.queue.put_nowait(data)
#
# def stop(self):
# self.fd.close()
#
# @property
# def queue(self):
# if not self.__queue:
# self.__queue = asyncio.Queue()
# return self.__queue
#
# @queue.setter
# def queue(self, queue):
# self.__queue = queue
class Communicator:
def __init__(self, loop, options):
self.logid = 'Communicator'
self.log = logging.getLogger(self.logid)
self.loop = loop
self.queuetask = None
self.queue = asyncio.Queue()
if options.rid:
self.nat = EcoNat(loop, server=options.server)
else:
self.nat = RadiusClient(loop, server=options.server, port=options.port, secret=options.secret.encode())
self.db = DbInfo(loop, user=options.user, password=options.password, dbname=options.dbname)
self.db.queue = self.queue
if options.rid:
self.ridtask = self.loop.create_task(self.rid())
else:
self.ridtask = None
@property
def queue(self):
if not self.__queue:
self.queue = asyncio.Queue()
return self.__queue
@queue.setter
def queue(self, queue):
self.__queue = queue
self.queuetask = self.loop.create_task(self.queueloop())
def stop(self):
if self.queuetask:
self.queuetask.cancel()
if self.ridtask:
self.ridtask.cancel()
if self.db:
self.db.stop()
if self.nat:
self.nat.stop()
async def rid(self):
while True:
try:
dbrid = await self.db.rid()
if dbrid == None:
self.log.info('no dbrid info from database')
continue
self.log.debug('dbrid %s', dbrid)
natrid = await self.nat.rid()
self.log.debug('natrid %s', natrid)
added, removed, modified = dict_compare(dbrid,natrid)
self.log.debug('added %s', added)
self.log.debug('removed %s', removed)
self.log.debug('modified %s', modified)
mlist = {}
if removed:
mlist['del'] = {'id': removed}
toadd = added.union(modified)
if toadd:
mlist['add'] = {'id': toadd}
if mlist:
await self.queue.put(mlist)
finally:
self.log.debug('sleep %d', 30)
await asyncio.sleep(30)
async def queueloop(self):
pool = {}
while True:
self.log.log(1, 'queueloop wait')
part = await self.queue.get()
self.log.debug('part %s', str(part))
try:
# get all data from queue:
if not isinstance(part, dict):
continue
for action, users in part.items():
if not isinstance(users, dict):
continue
for ids, values in users.items():
if isinstance(values, (int, str)):
values = [values]
pool.setdefault(action, {}).setdefault(ids, set()).update(values)
# get info from db:
if self.queue.empty():
self.log.debug('pool %s', str(pool))
action_users = {}
for action, users in pool.items():
for ids, values in users.items():
#userlist = None
#if action == 'add':
#userlist = await self.db.userlist(ids, values)
#elif action == 'del':
#if ids == 'id':
# userlist = {User(id=userid) for userid in values}
#else:
# userlist = await self.db.userlist(ids, values)
userlist = await self.db.userlist(self.nat.user, ids, values)
if userlist:
action_users.setdefault(action, []).extend(userlist)
pool.clear()
# send to activator
if action_users:
await self.nat.processusers(action_users)
except Exception as e:
self.log.exception('queueloop')
self.log.log(1, 'queueloop done')
def getoptions():
import argparse
parser = argparse.ArgumentParser(
description="{0} {1}".format(servicedesc, version))
parser.add_argument("-v", "--verbose",
dest="verbosity",
action="count",
help="print more diagnostic messages (option can be given multiple times)",
default=0
)
parser.add_argument("-l", "--log",
dest="logfile",
nargs="?",
help="log file, default: %(default)s, %(const)s if enabled",
const="/var/log/{0}/{0}.log".format(servicename)
)
parser.add_argument("-s", "--syslog",
dest="syslog",
action="store_true",
help="log to syslog (default off)",
default=False
)
parser.add_argument("-p", "--pid",
dest="pid",
nargs="?",
help="pid file, default: %(default)s, %(const)s if enabled",
const="/run/{0}/{0}.pid".format(servicename)
)
parser.add_argument("-f", "--foreground",
dest="foreground",
action="store_true",
help="stay in foreground (default off)",
default=False
)
parser.add_argument("--socket",
dest="pipe",
help="pipe file, default: %(default)s",
default="/var/run/{0}/{0}.socket".format(servicename)
)
group = parser.add_argument_group(
"econat server",
"econat server settings"
)
group.add_argument("--rid",
dest="rid",
action="store_true",
help="use RID insted of radius, default: use radius",
default=False
)
group.add_argument("--server",
dest="server",
help="econat server ip, default: %(default)s",
default="192.168.100.200"
)
group.add_argument("--port",
dest="port",
type=int,
help="econat server port, default: %(default)s",
default=1812
)
group.add_argument("--secret",
dest="secret",
help="econat radius coa_password, default: %(default)s",
default="radius"
)
group = parser.add_argument_group(
"database",
"Postgresql database connection settings"
)
group.add_argument("--user",
dest="user",
help="database user name, default: %(default)s",
default="radius"
)
group.add_argument("--password",
dest="password",
help="database password, default: %(default)s",
default="radius"
)
group.add_argument("--dbname",
dest="dbname",
help="database name to connect to, default: %(default)s",
default="radius"
)
return parser.parse_args()
def main():
'''start service as daemon'''
options = getoptions()
try:
setproctitle.setproctitle(servicename)
with daemon.DaemonContext(
pidfile = daemon.pidfile.PIDLockFile(options.pid) if options.pid else None,
signal_map = {signal.SIGTERM: lambda signum, stack_frame: sys.exit(0)},
detach_process = not options.foreground,
stdout = sys.stdout if options.foreground else None,
stderr = sys.stderr if options.foreground else None,
#uid = pwd.getpwnam(options.uid).pw_uid if options.uid else None,
#gid = grp.getgrnam(options.gid).gr_gid if options.gid else None,
files_preserve = [3] if 'LISTEN_FDNAMES' in os.environ else None,
) as context:
_main(options)
except lockfile.LockFailed as ex:
sys.stderr.write("daemonize error: {}'\n".format(ex))
def _main(options):
if options.verbosity > 3:
options.verbosity = 3
level = (
logging.WARNING,
logging.INFO,
logging.DEBUG,
logging.NOTSET,
)[options.verbosity]
logger = logging.getLogger('')
logger.addHandler(logging.NullHandler())
logger.setLevel(level)
logformat = '%(asctime)s %(levelname)s:%(name)s: %(message)s'
if options.logfile:
filelogger = logging.handlers.WatchedFileHandler(options.logfile)
filelogger.setFormatter(logging.Formatter(logformat))
logger.addHandler(filelogger)
if options.syslog:
syslogger = logging.handlers.SysLogHandler(address = '/dev/log', facility = logging.handlers.SysLogHandler.LOG_LOCAL0)
syslogger.setFormatter(logging.Formatter('%(name)s: %(message)s'))
logger.addHandler(syslogger)
if options.foreground:
conslogger = logging.StreamHandler()
conslogger.setFormatter(logging.Formatter(logformat))
logger.addHandler(conslogger)
sys.excepthook = lambda excType, excValue, traceback: logging.getLogger('exception').error("Uncaught exception", exc_info=(excType, excValue, traceback))
log = logging.getLogger(servicename)
log.info("starting %s version %s", servicename, version)
loop = asyncio.get_event_loop()
def callback(loop, exception):
log.error('ex: %s', exception)
loop.set_exception_handler(callback)
#fifo = UnixSocket(loop, file=options.pipe)
#fifo = FifoPipe(loop, file=options.pipe)
com = Communicator(loop, options)
#fifo.queue = com.queue
log.info('run loop')
try:
loop.run_forever()
except KeyboardInterrupt:
pass
finally:
com.stop()
#fifo.stop()
loop.stop()
loop.run_forever() # tasks cancel finishing
loop.close()
log.info("exit")
if __name__ == '__main__':
main()