add database, cli options, logging

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()

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"},

pyproject.toml

@@ -7,6 +7,7 @@ authors = ["Tomáš Mládek <t@mldk.cz>"]
 [tool.poetry.dependencies]
 python = "^3.7"
 routeros-api = "^0.17.0"
+click = "^7.1.2"
 
 [tool.poetry.dev-dependencies]