loadbalancer works

scripts/loadbalancer/cron.d/ofm_loadbalancer

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

scripts/loadbalancer/loadbalancer.py

@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
-
+import datetime
 import json
 
 import click
 import requests
+from dotenv import dotenv_values
+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
 
@@ -19,33 +21,62 @@ def cli():
 
 
 @cli.command()
-def run():
+def check():
     """
-    Runs load-balancing job (triggered by cron every minute)
+    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('/data/ofm/config/loadbalancer.json') as fp:
         c = json.load(fp)
-    # print(c)
+        # print(c)
 
     try:
         results_by_ip = {}
+        working_hosts = set()
 
         for area in AREAS:
-            for host_ip, host_ok in run_area(c, area).items():
+            for host_ip, host_is_ok in run_area(c, area).items():
                 results_by_ip.setdefault(host_ip, True)
-                results_by_ip[host_ip] &= host_ok
+                results_by_ip[host_ip] &= host_is_ok
 
-        for host_ip, host_ok in results_by_ip.items():
-            if not host_ok:
-                message = f'ERROR with host: {host_ip}'
+        for host_ip, host_is_ok in results_by_ip.items():
+            if not host_is_ok:
+                message = f'OFM ERROR with host: {host_ip}'
                 print(message)
                 telegram_send_message(message, c['telegram_token'], c['telegram_chat_id'])
+            else:
+                working_hosts.add(host_ip)
 
     except Exception as e:
-        message = f'ERROR with loadbalancer: {e}'
+        message = f'OFM ERROR with loadbalancer: {e}'
         print(message)
         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'])
+
+        update_records(c, working_hosts)
 
 
 def run_area(c, area):
@@ -53,7 +84,7 @@ def run_area(c, area):
 
     print(f'target version: {area}: {target_version}')
 
-    results = dict()
+    results = {}
 
     for host_ip in c['http_host_list']:
         try:
@@ -86,5 +117,32 @@ def get_target_version(area):
     return response.text.strip()
 
 
+def update_records(c, working_hosts):
+    config = dotenv_values('/data/ofm/config/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)
+
+    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,
+    )
+
+    set_records_round_robin(
+        zone_id=zone_id,
+        name=c['domain_cf'],
+        host_ip_set=working_hosts,
+        proxied=True,
+        comment='domain_cf',
+        cloudflare_api_token=cloudflare_api_token,
+    )
+
+
 if __name__ == '__main__':
     cli()

scripts/loadbalancer/loadbalancer_lib/cloudflare.py

@@ -0,0 +1,107 @@
+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,
+):
+    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
+
+    # 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
+
+
+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=dict(),
+    )
+    res.raise_for_status()
+    data = res.json()
+    assert data['success'] is True

scripts/loadbalancer/setup.py

@@ -5,6 +5,7 @@ requirements = [
     'click',
     'requests',
     'pycurl',
+    'python-dotenv',
 ]
 
 

ssh_lib/tasks.py

@@ -222,11 +222,13 @@ def setup_ledns_writer(c):
 
 
 def setup_loadbalancer(c):
+    domain_cf = dotenv_val('DOMAIN_CF').lower()
     domain_ledns = dotenv_val('DOMAIN_LEDNS').lower()
     http_host_list = [h.strip() for h in dotenv_val('HTTP_HOST_LIST').split(',') if h.strip()]
     assert (CONFIG_DIR / 'cloudflare.ini').exists()
 
     config = {
+        'domain_cf': domain_cf,
         'domain_ledns': domain_ledns,
         'http_host_list': http_host_list,
         'telegram_token': dotenv_val('TELEGRAM_TOKEN'),
@@ -234,7 +236,7 @@ def setup_loadbalancer(c):
     }
 
     config_str = json.dumps(config, indent=2, ensure_ascii=False)
-    print(config_str)
+    # print(config_str)
     put_str(c, f'{REMOTE_CONFIG}/loadbalancer.json', config_str)
 
     put(
@@ -253,3 +255,8 @@ def setup_loadbalancer(c):
     )
 
     c.sudo(f'{VENV_BIN}/pip install -e /data/ofm/loadbalancer')
+
+    c.sudo('mkdir -p /data/ofm/loadbalancer/logs')
+    put(c, SCRIPTS_DIR / 'loadbalancer' / 'cron.d' / 'ofm_loadbalancer', '/etc/cron.d/')
+
+    c.sudo('chown -R ofm:ofm /data/ofm/loadbalancer')