scripts -> modules

modules/loadbalancer/cron.d/ofm_loadbalancer

@@ -0,0 +1,8 @@
+# every minute
+
+# fix
+#* * * * * ofm     sudo /data/ofm/venv/bin/python -u /data/ofm/loadbalancer/loadbalancer.py fix >> /data/ofm/loadbalancer/logs/run.log 2>&1
+
+
+# check
+* * * * * ofm     sudo /data/ofm/venv/bin/python -u /data/ofm/loadbalancer/loadbalancer.py check >> /data/ofm/loadbalancer/logs/run.log 2>&1

modules/loadbalancer/loadbalancer.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+import datetime
+import json
+
+import click
+import requests
+from dotenv import dotenv_values
+from loadbalancer_lib import OFM_CONFIG_DIR
+from loadbalancer_lib.cloudflare import get_zone_id, set_records_round_robin
+from loadbalancer_lib.curl import pycurl_get, pycurl_status
+from loadbalancer_lib.telegram_ import telegram_send_message
+
+
+AREAS = ['planet', 'monaco']
+
+
+@click.group()
+def cli():
+    """
+    Manages load-balancing of Round-Robin DNS records
+    """
+
+
+@cli.command()
+def check():
+    """
+    Runs load-balancing check (triggered by cron every minute)
+    """
+
+    print(f'starting loadbalancer check at: {datetime.datetime.now(tz=datetime.timezone.utc)}')
+    check_or_fix(fix=False)
+
+
+@cli.command()
+def fix():
+    """
+    Fixes records based on check results
+    """
+
+    print(f'starting loadbalancer fix at: {datetime.datetime.now(tz=datetime.timezone.utc)}')
+    check_or_fix(fix=True)
+
+
+def check_or_fix(fix=False):
+    with open(OFM_CONFIG_DIR / 'loadbalancer.json') as fp:
+        c = json.load(fp)
+        # print(c)
+
+    if not c['http_host_list']:
+        telegram_send_message(
+            'OFM loadbalancer no hosts found on list, terminating',
+            c['telegram_token'],
+            c['telegram_chat_id'],
+        )
+        return
+
+    try:
+        results_by_ip = {}
+        working_hosts = set()
+
+        for area in AREAS:
+            results = run_area(c, area)
+            for host_ip, host_is_ok in results.items():
+                results_by_ip.setdefault(host_ip, True)
+                results_by_ip[host_ip] &= host_is_ok
+
+        for host_ip, host_is_ok in results_by_ip.items():
+            if not host_is_ok:
+                message = f'OFM loadbalancer ERROR with host: {host_ip}'
+                telegram_send_message(message, c['telegram_token'], c['telegram_chat_id'])
+            else:
+                working_hosts.add(host_ip)
+
+    except Exception as e:
+        message = f'OFM loadbalancer ERROR with loadbalancer: {e}'
+        telegram_send_message(message, c['telegram_token'], c['telegram_chat_id'])
+        return
+
+    print(f'working hosts: {sorted(working_hosts)}')
+
+    if fix:
+        # if no hosts are detected working, probably a bug in this script
+        # fail-safe to include all hosts
+        if not working_hosts:
+            working_hosts = set(c['http_host_list'])
+
+            message = 'OFM loadbalancer FIX found no working hosts, reverting to full list!'
+            telegram_send_message(message, c['telegram_token'], c['telegram_chat_id'])
+
+        updated = update_records(c, working_hosts)
+        if updated:
+            message = f'OFM loadbalancer FIX modified records, new records: {working_hosts}'
+            telegram_send_message(message, c['telegram_token'], c['telegram_chat_id'])
+
+
+def run_area(c, area):
+    target_version = get_target_version(area)
+
+    print(f'target version: {area}: {target_version}')
+
+    results = {}
+
+    for host_ip in c['http_host_list']:
+        try:
+            check_host(c['domain_ledns'], host_ip, area, target_version)
+            results[host_ip] = True
+        except Exception as e:
+            results[host_ip] = False
+            print(e)
+
+    return results
+
+
+def check_host(domain, host_ip, area, version):
+    # check TileJSON first
+    url = f'https://{domain}/{area}'
+    tilejson_str = pycurl_get(url, domain, host_ip)
+    tilejson = json.loads(tilejson_str)
+    tiles_url = tilejson['tiles'][0]
+    version_in_tilejson = tiles_url.split('/')[4]
+    assert version_in_tilejson == version
+
+    # check actual vector tile
+    url = f'https://{domain}/{area}/{version}/14/8529/5975.pbf'
+    assert pycurl_status(url, domain, host_ip) == 200
+
+    # check style
+    url = f'https://{domain}/styles/bright'
+    assert pycurl_status(url, domain, host_ip) == 200
+
+
+def get_target_version(area):
+    url = f'https://assets.openfreemap.com/versions/deployed_{area}.txt'
+    response = requests.get(url)
+    response.raise_for_status()
+    return response.text.strip()
+
+
+def update_records(c, working_hosts) -> bool:
+    config = dotenv_values(OFM_CONFIG_DIR / 'cloudflare.ini')
+    cloudflare_api_token = config['dns_cloudflare_api_token']
+
+    domain = '.'.join(c['domain_ledns'].split('.')[-2:])
+    zone_id = get_zone_id(domain, cloudflare_api_token=cloudflare_api_token)
+
+    updated = False
+
+    updated |= set_records_round_robin(
+        zone_id=zone_id,
+        name=c['domain_ledns'],
+        host_ip_set=working_hosts,
+        proxied=False,
+        ttl=300,
+        comment='domain_ledns',
+        cloudflare_api_token=cloudflare_api_token,
+    )
+
+    return updated
+
+
+if __name__ == '__main__':
+    cli()

modules/loadbalancer/loadbalancer_lib/__init__.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+
+
+if Path('/data/ofm/config').exists():
+    OFM_CONFIG_DIR = Path('/data/ofm/config')
+else:
+    OFM_CONFIG_DIR = Path(__file__).parent.parent.parent.parent / 'config'
+
+assert OFM_CONFIG_DIR.exists()

modules/loadbalancer/loadbalancer_lib/cloudflare.py

@@ -0,0 +1,109 @@
+from pprint import pprint
+
+import requests
+
+
+# docs: https://api.cloudflare.com/
+
+
+def cloudflare_get(path: str, params: dict, cloudflare_api_token: str):
+    headers = {'Authorization': f'Bearer {cloudflare_api_token}'}
+    res = requests.get(
+        f'https://api.cloudflare.com/client/v4{path}', headers=headers, params=params
+    )
+    res.raise_for_status()
+    data = res.json()
+    assert data['success'] is True
+    return data
+
+
+def get_zone_id(domain, cloudflare_api_token: str):
+    data = cloudflare_get(
+        '/zones', params=dict(name=domain), cloudflare_api_token=cloudflare_api_token
+    )
+    assert len(data['result']) == 1
+    zone_info = data['result'][0]
+    return zone_info['id']
+
+
+def get_dns_records_round_robin(zone_id, cloudflare_api_token: str) -> dict:
+    data = cloudflare_get(
+        f'/zones/{zone_id}/dns_records',
+        params=dict(per_page=5000),
+        cloudflare_api_token=cloudflare_api_token,
+    )
+    records = data['result']
+
+    data = {}
+
+    for r in records:
+        if r['type'] != 'A':
+            continue
+
+        data.setdefault(r['name'], [])
+        data[r['name']].append(dict(content=r['content'], id=r['id']))
+
+    return data
+
+
+def set_records_round_robin(
+    zone_id,
+    *,
+    name: str,
+    host_ip_set: set,
+    ttl: int = 1,
+    proxied: bool,
+    comment: str = None,
+    cloudflare_api_token: str,
+) -> bool:
+    headers = {'Authorization': f'Bearer {cloudflare_api_token}'}
+
+    dns_records = get_dns_records_round_robin(zone_id, cloudflare_api_token=cloudflare_api_token)
+    current_records = dns_records.get(name, [])
+
+    current_ips = {r['content'] for r in current_records}
+    if current_ips == host_ip_set:
+        print(f'No need to update records: {name} currently set: {sorted(current_ips)}')
+        return False
+
+    # changing records
+
+    # delete all current records first
+    for r in current_records:
+        delete_record(zone_id, id_=r['id'], cloudflare_api_token=cloudflare_api_token)
+
+    # create new records
+    for ip in host_ip_set:
+        print(f'Creating record: {name} {ip}')
+        json_data = dict(
+            type='A',
+            name=name,
+            content=ip,
+            ttl=ttl,
+            proxied=proxied,
+            comment=comment,
+        )
+        res = requests.post(
+            f'https://api.cloudflare.com/client/v4/zones/{zone_id}/dns_records',
+            headers=headers,
+            json=json_data,
+        )
+        res.raise_for_status()
+        data = res.json()
+        assert data['success'] is True
+
+    return True
+
+
+def delete_record(zone_id, *, id_: str, cloudflare_api_token: str):
+    headers = {'Authorization': f'Bearer {cloudflare_api_token}'}
+
+    print(f'Deleting record: {id_}')
+    res = requests.delete(
+        f'https://api.cloudflare.com/client/v4/zones/{zone_id}/dns_records/{id_}',
+        headers=headers,
+        json={},
+    )
+    res.raise_for_status()
+    data = res.json()
+    assert data['success'] is True

modules/loadbalancer/loadbalancer_lib/curl.py

@@ -0,0 +1,54 @@
+from io import BytesIO
+from pathlib import Path
+
+import pycurl
+
+
+def pycurl_status(url, domain, host_ip):
+    """
+    Uses pycurl to make a HTTPS HEAD request using custom resolving,
+    checks if the status code is 200
+    """
+
+    c = pycurl.Curl()
+    c.setopt(c.URL, url)
+
+    # linux needs CA certs specified manually
+    if Path('/etc/ssl/certs/ca-certificates.crt').exists():
+        c.setopt(c.CAINFO, '/etc/ssl/certs/ca-certificates.crt')
+
+    c.setopt(c.RESOLVE, [f'{domain}:443:{host_ip}'])
+    c.setopt(c.NOBODY, True)
+    c.setopt(c.TIMEOUT, 5)
+    c.perform()
+    status_code = c.getinfo(c.RESPONSE_CODE)
+    c.close()
+
+    return status_code
+
+
+def pycurl_get(url, domain, host_ip):
+    """
+    Uses pycurl to make a HTTPS GET request using custom resolving,
+    checks if the status code is 200, and returns the content.
+    """
+
+    buffer = BytesIO()
+    c = pycurl.Curl()
+    c.setopt(c.URL, url)
+
+    # linux needs CA certs specified manually
+    if Path('/etc/ssl/certs/ca-certificates.crt').exists():
+        c.setopt(c.CAINFO, '/etc/ssl/certs/ca-certificates.crt')
+
+    c.setopt(c.RESOLVE, [f'{domain}:443:{host_ip}'])
+    c.setopt(c.WRITEDATA, buffer)
+    c.setopt(c.TIMEOUT, 5)
+    c.perform()
+    status_code = c.getinfo(c.RESPONSE_CODE)
+    c.close()
+
+    if status_code != 200:
+        raise ValueError(f'status code: {status_code}')
+
+    return buffer.getvalue().decode('utf8')

modules/loadbalancer/loadbalancer_lib/telegram_.py

@@ -0,0 +1,16 @@
+import requests
+
+
+def telegram_send_message(message, bot_token, chat_id):
+    print(message)
+
+    url = f'https://api.telegram.org/bot{bot_token}/sendMessage'
+
+    payload = {'chat_id': chat_id, 'text': message}
+
+    response = requests.post(url, data=payload)
+
+    if response.status_code == 200:
+        print('  Message sent successfully!')
+    else:
+        print('  Failed to send message:', response.text)

modules/loadbalancer/setup.py

@@ -0,0 +1,16 @@
+from setuptools import find_packages, setup
+
+
+requirements = [
+    'click',
+    'requests',
+    'pycurl',
+    'python-dotenv',
+]
+
+
+setup(
+    python_requires='>=3.10',
+    install_requires=requirements,
+    packages=find_packages(),
+)