loadbalancer works

modules/loadbalancer/loadbalancer.py

@@ -1,16 +1,12 @@
 #!/usr/bin/env python3
-import json
+
 from datetime import datetime, timezone
 
 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.telegram_ import telegram_send_message
+from loadbalancer_lib.loadbalance import check_or_fix
 
 
-AREAS = ['planet', 'monaco']
+now = datetime.now(timezone.utc)
 
 
 @click.group()
@@ -23,121 +19,22 @@ def cli():
 @cli.command()
 def check():
     """
-    Runs load-balancing check (triggered by cron every minute)
+    Runs load-balancing check
     """
 
-    print(f'starting loadbalancer check at: {datetime.now(timezone.utc)}')
+    print(f'---\n{now}\nStarting check')
     check_or_fix(fix=False)
 
 
 @cli.command()
 def fix():
     """
-    Fixes records based on check results
+    Runs check and fixes records based on check results
     """
 
-    print(f'starting loadbalancer fix at: {datetime.now(timezone.utc)}')
+    print(f'---\n{now}\nStarting fix')
     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 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

@@ -1,9 +0,0 @@
-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

@@ -1,5 +1,3 @@
-from pprint import pprint
-
 import requests
 
 
@@ -62,6 +60,7 @@ def set_records_round_robin(
     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

modules/loadbalancer/loadbalancer_lib/config.py

@@ -0,0 +1,29 @@
+import json
+from pathlib import Path
+
+from dotenv import dotenv_values
+
+
+class Configuration:
+    areas = ['planet', 'monaco']
+
+    if Path('/data/ofm').exists():
+        ofm_config_dir = Path('/data/ofm/config')
+    else:
+        repo_root = Path(__file__).parent.parent.parent.parent
+        ofm_config_dir = repo_root / 'config'
+
+    ofm_config = json.loads((ofm_config_dir / 'config.json').read_text())
+
+    http_host_list = ofm_config['http_host_list']
+    telegram_token = ofm_config['telegram_token']
+    telegram_chat_id = ofm_config['telegram_chat_id']
+
+    domain_ledns = ofm_config['domain_ledns']
+    domain_root = '.'.join(domain_ledns.split('.')[-2:])
+
+    cloudflare_ini = dotenv_values(ofm_config_dir / 'cloudflare.ini')
+    cloudflare_api_token = cloudflare_ini['dns_cloudflare_api_token']
+
+
+config = Configuration()

modules/loadbalancer/loadbalancer_lib/loadbalance.py

@@ -0,0 +1,90 @@
+from http_host_lib.shared import get_deployed_version
+
+from loadbalancer_lib.cloudflare import get_zone_id, set_records_round_robin
+from loadbalancer_lib.config import config
+from loadbalancer_lib.shared import check_host_latest
+from loadbalancer_lib.telegram_ import telegram_send_message
+
+
+def check_or_fix(fix=False):
+    if not config.http_host_list:
+        telegram_quick(
+            'OFM loadbalancer no hosts found on list, terminating',
+        )
+        return
+
+    try:
+        results_by_ip = {}
+        working_hosts = set()
+
+        for area in config.areas:
+            results = run_area(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:
+                telegram_quick(f'OFM loadbalancer ERROR with host: {host_ip}')
+            else:
+                working_hosts.add(host_ip)
+
+    except Exception as e:
+        telegram_quick(f'OFM loadbalancer ERROR with loadbalancer: {e}')
+        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(config.http_host_list)
+            telegram_quick('OFM loadbalancer FIX found no working hosts, reverting to full list!')
+
+        updated = update_records(working_hosts)
+        if updated:
+            telegram_quick(f'OFM loadbalancer FIX modified records, new records: {working_hosts}')
+
+
+def run_area(area):
+    version = get_deployed_version(area)
+    if not version:
+        print(f'  deployed version not found: {area}')
+        return
+
+    print(f'  deployed version {area}: {version}')
+
+    results = {}
+
+    for host_ip in config.http_host_list:
+        try:
+            check_host_latest(config.domain_ledns, host_ip, area, version)
+            results[host_ip] = True
+        except Exception as e:
+            results[host_ip] = False
+            print(e)
+
+    return results
+
+
+def update_records(working_hosts) -> bool:
+    zone_id = get_zone_id(config.domain_root, cloudflare_api_token=config.cloudflare_api_token)
+
+    updated = False
+
+    updated |= set_records_round_robin(
+        zone_id=zone_id,
+        name=config.domain_ledns,
+        host_ip_set=working_hosts,
+        proxied=False,
+        ttl=300,
+        comment='domain_ledns',
+        cloudflare_api_token=config.cloudflare_api_token,
+    )
+
+    return updated
+
+
+def telegram_quick(message):
+    telegram_send_message(message, config.telegram_token, config.telegram_chat_id)