From 29399ab34607a03a1ed4421f838a5afb629f34c7 Mon Sep 17 00:00:00 2001 From: Zsolt Ero Date: Mon, 10 Jun 2024 01:16:31 +0200 Subject: [PATCH] loadbalancer works --- scripts/loadbalancer/cron.d/ofm_loadbalancer | 4 + scripts/loadbalancer/loadbalancer.py | 80 +++++++++++-- .../loadbalancer_lib/cloudflare.py | 107 ++++++++++++++++++ scripts/loadbalancer/setup.py | 1 + ssh_lib/tasks.py | 9 +- 5 files changed, 189 insertions(+), 12 deletions(-) create mode 100644 scripts/loadbalancer/cron.d/ofm_loadbalancer create mode 100644 scripts/loadbalancer/loadbalancer_lib/cloudflare.py diff --git a/scripts/loadbalancer/cron.d/ofm_loadbalancer b/scripts/loadbalancer/cron.d/ofm_loadbalancer new file mode 100644 index 0000000..3adca38 --- /dev/null +++ b/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 + + diff --git a/scripts/loadbalancer/loadbalancer.py b/scripts/loadbalancer/loadbalancer.py index 8df1b06..9d86360 100755 --- a/scripts/loadbalancer/loadbalancer.py +++ b/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() diff --git a/scripts/loadbalancer/loadbalancer_lib/cloudflare.py b/scripts/loadbalancer/loadbalancer_lib/cloudflare.py new file mode 100644 index 0000000..17f171c --- /dev/null +++ b/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 diff --git a/scripts/loadbalancer/setup.py b/scripts/loadbalancer/setup.py index 8734056..84f9323 100644 --- a/scripts/loadbalancer/setup.py +++ b/scripts/loadbalancer/setup.py @@ -5,6 +5,7 @@ requirements = [ 'click', 'requests', 'pycurl', + 'python-dotenv', ] diff --git a/ssh_lib/tasks.py b/ssh_lib/tasks.py index f8a9ef0..3b8fe6c 100644 --- a/ssh_lib/tasks.py +++ b/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')