From 11846f1b4f49aa9095f1c10b019657d27103db48 Mon Sep 17 00:00:00 2001
From: redxef <federanko.manuel@gmail.com>
Date: Wed, 5 Jun 2024 23:46:04 +0200
Subject: [PATCH] Allow multiple subdomains.

---
 main.py | 74 +++++++++++++++++++++++++++++++++------------------------
 1 file changed, 43 insertions(+), 31 deletions(-)

diff --git a/main.py b/main.py
index 788b9fb..170662f 100755
--- a/main.py
+++ b/main.py
@@ -3,6 +3,7 @@
 import itertools
 import os
 import sys
+import typing
 import logging
 
 import yaml
@@ -15,6 +16,12 @@ from schema import And, Optional, Use, Schema
 
 logger = logging.getLogger('ovhdns')
 
+def str_or_list(i: typing.Union[str, typing.List[str]]) -> typing.List[str]:
+    if isinstance(i, str):
+        return [i]
+    return i
+
+
 config_schema = Schema({
     'ovh': {
         'application_key': Use(str),
@@ -23,7 +30,7 @@ config_schema = Schema({
         'endpoint': Use(str),
     },
     'zone': Use(str),
-    Optional('subdomain', default=''): Use(str),
+    Optional('subdomain', default=''): Use(str_or_list),
     Optional('field_type', default='A'): And(Use(str), Use(str.upper), dns.rdatatype.RdataType.make),
     Optional('logging', default={'level': ':WARNING,ovhdns:INFO'}): {
         Optional('level', default=':WARNING,ovhdns:INFO'): Use(str),
@@ -108,44 +115,49 @@ def main(config, level, force_refresh):
     logger.debug('Found records %d:', len(records))
     for r in records:
         logger.debug(' - %s', repr(r))
-    records = [
-        r
-        for r in records
-        if r['fieldType'] == config['field_type']
-        and r['zone'] == config['zone']
-        and r['subDomain'] == config['subdomain']
-    ]
-    logger.info('found matching records: %s', records)
-    if len(records) > 1:
-        logger.error('found more than one record, don\'t know which to pick')
-        sys.exit(1)
+    logging.info('matching against subdomains %s', config['subdomain'])
+    records_dict = {}
+    for subdomain in config['subdomain']:
+        records_dict[subdomain] = [
+            r
+            for r in records
+            if r['fieldType'] == config['field_type']
+            and r['zone'] == config['zone']
+            and r['subDomain'] == subdomain
+        ]
+    logger.info('found matching records: %s', records_dict)
+    for v in records_dict.values():
+        if len(v) > 1:
+            logger.error('found more than one record, don\'t know which to pick')
+            sys.exit(1)
 
     my_ips = get_my_ips()
     need_refresh = False
-    if len(records) == 1:
-        logger.info('record already present, updating if needed')
-        if my_ips[0] == records[0]['target']:
-            logger.info('ip is up-to-date, skipping update')
+    for subdomain, records in records_dict.items():
+        if len(records) == 1:
+            logger.info('record already present, updating if needed')
+            if my_ips[0] == records[0]['target']:
+                logger.info('ip is up-to-date, skipping update')
+            else:
+                need_refresh = True
+                ovh_client.put(
+                    f'/domain/zone/{config["zone"]}/record/{records[0]["id"]}',
+                    target=my_ips[0],
+                    ttl=config['ttl'],
+                )
+                logger.info('updated record, new ip is %s', my_ips[0])
+
         else:
+            logger.info('record not already present, creating new one')
             need_refresh = True
-            ovh_client.put(
-                f'/domain/zone/{config["zone"]}/record/{records[0]["id"]}',
+            ovh_client.post(
+                f'/domain/zone/{config["zone"]}/record/',
+                subDomain=subdomain,
                 target=my_ips[0],
+                fieldType=config['field_type'],
                 ttl=config['ttl'],
             )
-            logger.info('updated record, new ip is %s', my_ips[0])
-
-    else:
-        logger.info('record not already present, creating new one')
-        need_refresh = True
-        ovh_client.post(
-            f'/domain/zone/{config["zone"]}/record/',
-            subDomain=config['subdomain'],
-            target=my_ips[0],
-            fieldType=config['field_type'],
-            ttl=config['ttl'],
-        )
-        logger.info('created new record, ip is %s', my_ips[0])
+            logger.info('created new record, ip is %s', my_ips[0])
 
     if need_refresh or force_refresh:
         ovh_client.post(f'/domain/zone/{config["zone"]}/refresh')