#!/usr/bin/env python3
"""Conservative XML edits preserving unrelated Wazuh configuration/comments."""
from pathlib import Path
import re
import sys
from xml.dom import minidom

LOG = '/var/log/wazuh-docker-security/events.jsonl'


def parse_fragment(text):
    if '<!DOCTYPE' in text.upper() or '<!ENTITY' in text.upper():
        raise ValueError('XML entities/DOCTYPE are unsupported')
    text = re.sub(r'^\s*<\?xml[^?]*\?>', '', text, count=1)
    return minidom.parseString('<configuration_fragments>' + text + '</configuration_fragments>')


def element_children(parent, tag):
    return [n for n in parent.childNodes if n.nodeType == n.ELEMENT_NODE and n.tagName == tag]


def value(node):
    return ''.join(n.data for n in node.childNodes if n.nodeType == n.TEXT_NODE).strip()


def set_child(document, parent, tag, text):
    nodes = element_children(parent, tag)
    if len(nodes) > 1:
        raise ValueError('ambiguous duplicate ' + tag)
    node = nodes[0] if nodes else document.createElement(tag)
    if not nodes:
        parent.appendChild(node)
    for child in list(node.childNodes):
        node.removeChild(child)
    node.appendChild(document.createTextNode(text))


def configure(text, listener, scans):
    doc = parse_fragment(text)
    roots = element_children(doc.documentElement, 'ossec_config')
    if not roots:
        raise ValueError('ossec_config missing')
    if listener:
        matches = [n for root in roots for n in element_children(root, 'wodle')
                   if n.getAttribute('name') == 'docker-listener']
        if len(matches) > 1:
            raise ValueError('multiple docker-listener blocks; resolve manually')
        node = matches[0] if matches else doc.createElement('wodle')
        if not matches:
            node.setAttribute('name', 'docker-listener')
            roots[0].appendChild(node)
        set_child(doc, node, 'disabled', 'no')
    if scans:
        matches = [n for root in roots for n in element_children(root, 'localfile')
                   if any(value(x) == LOG for x in element_children(n, 'location'))]
        if len(matches) > 1:
            raise ValueError('multiple managed scan log blocks')
        node = matches[0] if matches else doc.createElement('localfile')
        if not matches:
            roots[0].appendChild(doc.createComment(' Managed by wazuh-docker-security '))
            roots[0].appendChild(node)
        set_child(doc, node, 'location', LOG)
        set_child(doc, node, 'log_format', 'json')
        set_child(doc, node, 'only-future-events', 'no')
    return ''.join(n.toxml() for n in doc.documentElement.childNodes).rstrip() + '\n'


def check_rule_collisions(rules_dir, target, new_rules):
    wanted = set(re.findall(r'<rule\s+id="(\d+)"', new_rules))
    for directory in (Path(rules_dir), Path(rules_dir).parents[1] / 'ruleset/rules'):
        if not directory.exists():
            continue
        for path in directory.glob('*.xml'):
            if path.resolve() == Path(target).resolve():
                continue
            existing = set(re.findall(r'<rule\b[^>]*\bid\s*=[\'"](\d+)[\'"]', path.read_text()))
            if wanted & existing:
                raise ValueError('rule IDs already used in ' + str(path))


def main():
    mode = sys.argv[1]
    if mode == 'agent':
        source, output = map(Path, sys.argv[2:4])
        if source.is_symlink() or not source.is_file():
            raise ValueError('ossec.conf must be a regular file')
        output.write_text(configure(source.read_text(), sys.argv[4] == '1', sys.argv[5] == '1'))
    elif mode == 'rules':
        source, target = map(Path, sys.argv[2:4])
        new_rules = source.read_text()
        parse_fragment(new_rules)
        if target.is_symlink():
            raise ValueError('rule file is a symlink')
        if target.exists() and 'Managed by wazuh-docker-security' not in target.read_text():
            raise ValueError('unmanaged rule file already exists')
        check_rule_collisions(target.parent, target, new_rules)
    else:
        raise ValueError('unknown mode')


if __name__ == '__main__':
    main()