From be4273596b8b859895f48abfec1ef2d1d1719c7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Ml=C3=A1dek?= Date: Tue, 9 Feb 2021 20:44:19 +0100 Subject: [PATCH] add database, cli options, logging --- generate.py | 114 +++++++++++++++++++++++++++++++++++++++++-------- poetry.lock | 42 +++++++++++++++++- pyproject.toml | 1 + 3 files changed, 138 insertions(+), 19 deletions(-) diff --git a/generate.py b/generate.py index d4bded8..63e631a 100644 --- a/generate.py +++ b/generate.py @@ -1,21 +1,101 @@ -import routeros_api -import os -import sys import csv +import logging +import os +import sqlite3 +from dataclasses import dataclass +from datetime import datetime +from time import sleep +from typing import List, Optional -connection = routeros_api.RouterOsApiPool( - os.getenv("ROUTER_IP", "192.168.42.1"), - username='admin', - password=os.getenv("ROUTER_PASSWORD"), - plaintext_login=True -) -api = connection.get_api() +import click +import routeros_api -registered = api.get_resource('/caps-man/registration-table').call('print') -leases = api.get_resource('/ip/dhcp-server/lease').call('print') -writer = csv.writer(sys.stdout) -for client in registered: - if os.getenv("SSID", "anabasis") in client['ssid']: - lease = next(l for l in leases if l.get('active-mac-address') == client['mac-address']) - writer.writerow((lease['active-address'], lease['active-mac-address'], lease.get('host-name', "???"))) +def _get_db(filepath: str): + logging.debug(f"Opening database: {filepath}") + connection = sqlite3.connect(filepath) + + c = connection.cursor() + c.execute( + '''CREATE TABLE IF NOT EXISTS spottings ( + id INTEGER PRIMARY KEY, + ts INTEGER, + mac TEXT, + hostname TEXT, + ip TEXT + )''' + ) + connection.commit() + + return connection + + +@dataclass +class Lease: + mac: str + hostname: Optional[str] + ip: str + + +@click.command() +@click.option('--address', default=os.getenv("ROUTER_IP", "192.168.42.1"), help='Address of the router.') +@click.option('--period', default=60, help='How often to check for clients (in seconds).') +@click.option('--ssid', default=os.getenv("SSID", "anabasis"), help='Limit clients to SSID containing this string.') +@click.option('-o', '--output', default="clients.lst", help="Output CSV file.") +def run_forever(address: str, period: int, ssid: str, output: str): + logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - [%(levelname)s] %(message)s') + db = _get_db("clients.sqlite3") + + connection = routeros_api.RouterOsApiPool( + address, + username='admin', + password=os.environ["ROUTER_PASSWORD"], + plaintext_login=True + ) + + while True: + logging.info(f"Querying router at {address}...") + api = connection.get_api() + + currently_registered = api.get_resource('/caps-man/registration-table').call('print') + logging.debug(f"Got {len(currently_registered)} registered clients.") + dhcp_leases = api.get_resource('/ip/dhcp-server/lease').call('print') + logging.debug(f"Got {len(dhcp_leases)} DHCP leases.") + + registered_leases: List[Lease] = [] + for client in filter(lambda c: ssid in c['ssid'], currently_registered): + try: + lease = next( + lease for lease in dhcp_leases if lease.get('active-mac-address') == client['mac-address'] + ) + except StopIteration: + continue + registered_leases.append( + Lease(ip=lease['active-address'], mac=lease['active-mac-address'], hostname=lease.get('host-name')) + ) + logging.info(f"Found {len(registered_leases)} registered leases.") + logging.debug(", ".join([str(lease) for lease in registered_leases])) + + now = datetime.now() + timestamp = int(now.timestamp()) + + logging.debug("Logging into the database...") + cur = db.cursor() + cur.executemany("INSERT INTO spottings (ts, mac, hostname, ip) VALUES (?,?,?,?)", [ + (timestamp, lease.mac, lease.hostname, lease.ip) for lease in registered_leases + ]) + db.commit() + + logging.debug("Outputting CSV file...") + with open(output, 'w') as file: + file.write(f"{now}\n") + writer = csv.writer(file) + for lease in registered_leases: + writer.writerow((lease.ip, lease.mac, lease.hostname or "???")) + + logging.info(f"Sleeping for {period} seconds.") + sleep(period) + + +if __name__ == '__main__': + run_forever() diff --git a/poetry.lock b/poetry.lock index 56aab1c..67d00d3 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,3 +1,30 @@ +[[package]] +name = "argparse" +version = "1.4.0" +description = "Python command-line parsing library" +category = "main" +optional = false +python-versions = "*" + +[[package]] +name = "caribou" +version = "0.2.1" +description = "python migrations for sqlite databases" +category = "main" +optional = false +python-versions = "*" + +[package.dependencies] +argparse = ">=1.0.0" + +[[package]] +name = "click" +version = "7.1.2" +description = "Composable command line interface toolkit" +category = "main" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" + [[package]] name = "routeros-api" version = "0.17.0" @@ -18,11 +45,22 @@ optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" [metadata] -lock-version = "1.0" +lock-version = "1.1" python-versions = "^3.7" -content-hash = "3ab50217eea285c11ccd2f1b4ab50a57d19cdf03f1c8e71e6bfd63e3e225cb59" +content-hash = "1c76be1a79a755045b6946b321f49d54220a1f6061a643b83a305704f976fc61" [metadata.files] +argparse = [ + {file = "argparse-1.4.0-py2.py3-none-any.whl", hash = "sha256:c31647edb69fd3d465a847ea3157d37bed1f95f19760b11a47aa91c04b666314"}, + {file = "argparse-1.4.0.tar.gz", hash = "sha256:62b089a55be1d8949cd2bc7e0df0bddb9e028faefc8c32038cc84862aefdd6e4"}, +] +caribou = [ + {file = "caribou-0.2.1.tar.gz", hash = "sha256:2738a8404cb80c9b6209e5af29481d0aded9009a53265b45accf6a4ba04c2cbb"}, +] +click = [ + {file = "click-7.1.2-py2.py3-none-any.whl", hash = "sha256:dacca89f4bfadd5de3d7489b7c8a566eee0d3676333fbb50030263894c38c0dc"}, + {file = "click-7.1.2.tar.gz", hash = "sha256:d2b5255c7c6349bc1bd1e59e08cd12acbbd63ce649f2588755783aa94dfb6b1a"}, +] routeros-api = [ {file = "RouterOS-api-0.17.0.tar.gz", hash = "sha256:1b9898460ecc4667b54e477d495b74c2f24ae0aac4c90dd0e62f23ec7eae8252"}, {file = "RouterOS_api-0.17.0-py2.py3-none-any.whl", hash = "sha256:bf38da94a570875eaa87ff537558f765a4697dbce1a9753070194b687f441bf0"}, diff --git a/pyproject.toml b/pyproject.toml index 01df22c..f4e359a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,7 @@ authors = ["Tomáš Mládek "] [tool.poetry.dependencies] python = "^3.7" routeros-api = "^0.17.0" +click = "^7.1.2" [tool.poetry.dev-dependencies]