#!/usr/bin/env perl
#
# mdns_host_seed.pl - Collect mDNS A records into the host source database.
#

use strict;
use warnings;

use Cwd qw(abs_path);
use Fcntl qw(:flock);
use File::Basename qw(dirname);
use File::Path qw(make_path);
use Getopt::Long qw(GetOptions);
use IO::Socket::INET;
use POSIX qw(strftime);
use Socket qw(IPPROTO_IP IP_ADD_MEMBERSHIP SO_REUSEADDR inet_aton inet_ntoa sockaddr_in);

my $script_dir = dirname(abs_path($0));
my $project_dir = dirname($script_dir);

my %opt = (
    database => $ENV{HOST_MANAGER_MDNS_DB} || "$project_dir/var/mdns-observations.yaml",
    bind => $ENV{HOST_MANAGER_MDNS_BIND} || '0.0.0.0',
    group => $ENV{HOST_MANAGER_MDNS_GROUP} || '224.0.0.251',
    port => $ENV{HOST_MANAGER_MDNS_PORT} || 5353,
    once => 0,
    timeout => 0,
    dry_run => 0,
    verbose => 0,
);

GetOptions(
    'database|db=s' => \$opt{database},
    'bind=s' => \$opt{bind},
    'group=s' => \$opt{group},
    'port=i' => \$opt{port},
    'once' => \$opt{once},
    'timeout=i' => \$opt{timeout},
    'dry-run' => \$opt{dry_run},
    'verbose' => \$opt{verbose},
    'help|h' => sub { usage(); exit 0; },
) or do {
    usage();
    exit 2;
};

my $socket = IO::Socket::INET->new(
    LocalAddr => $opt{bind},
    LocalPort => $opt{port},
    Proto => 'udp',
    ReuseAddr => 1,
) or die "Cannot listen for mDNS on $opt{bind}:$opt{port}: $!\n";

setsockopt($socket, Socket::SOL_SOCKET(), SO_REUSEADDR, pack('i', 1));
setsockopt($socket, IPPROTO_IP, IP_ADD_MEMBERSHIP, inet_aton($opt{group}) . inet_aton($opt{bind}))
    or die "Cannot join mDNS multicast group $opt{group}: $!\n";

print "mdns host seed listening on udp://$opt{bind}:$opt{port} group $opt{group}\n" if $opt{verbose};
print "mDNS source database: $opt{database}\n" if $opt{verbose};

my $deadline = $opt{timeout} ? time() + $opt{timeout} : 0;
while (1) {
    my $timeout = undef;
    if ($deadline) {
        $timeout = $deadline - time();
        last if $timeout <= 0;
    }

    my $rin = '';
    vec($rin, fileno($socket), 1) = 1;
    my $ready = select(my $rout = $rin, undef, undef, $timeout);
    last if !defined($ready) || $ready <= 0;

    my $peer = recv($socket, my $packet, 9000, 0);
    next unless defined $peer && length $packet;

    my ($peer_port, $peer_addr) = sockaddr_in($peer);
    my $peer_ip = inet_ntoa($peer_addr);
    my @records = grep { record_is_usable($_) } parse_mdns_packet($packet);
    next unless @records;

    my $changed = seed_observations(\%opt, \@records, $peer_ip);
    print "stored " . scalar(@$changed) . " observation change(s)\n" if $opt{verbose} && @$changed;
    last if $opt{once} && @$changed;
}

exit 0;

sub usage {
    print <<"EOF";
Usage: perl scripts/mdns_host_seed.pl [options]

Options:
  --database path  mDNS source database. Defaults to var/mdns-observations.yaml.
  --bind addr      Local bind address. Defaults to 0.0.0.0.
  --group addr     Multicast group. Defaults to 224.0.0.251.
  --port n         UDP port. Defaults to 5353.
  --once           Exit after the first stored observation.
  --timeout n      Exit after n seconds.
  --dry-run        Print proposed observation changes without writing.
  --verbose        Print listener and change details.

Only A records for private/link-local .local names are collected. hosts.yaml is
not modified here; it remains a generated output fed by source databases.
EOF
}

sub parse_mdns_packet {
    my ($packet) = @_;
    return () if length($packet) < 12;
    my ($id, $flags, $qd, $an, $ns, $ar) = unpack('n6', substr($packet, 0, 12));
    my $offset = 12;

    for (1 .. $qd) {
        my ($name, $next) = read_dns_name($packet, $offset);
        last unless defined $name && $next + 4 <= length($packet);
        $offset = $next + 4;
    }

    my @records;
    for (1 .. ($an + $ns + $ar)) {
        my ($name, $next) = read_dns_name($packet, $offset);
        last unless defined $name && $next + 10 <= length($packet);
        my ($type, $class, $ttl, $rdlen) = unpack('nnNn', substr($packet, $next, 10));
        my $rdata_offset = $next + 10;
        last if $rdata_offset + $rdlen > length($packet);
        if ($type == 1 && $rdlen == 4) {
            push @records, {
                name => lc($name),
                ip => inet_ntoa(substr($packet, $rdata_offset, 4)),
                ttl => $ttl,
            };
        }
        $offset = $rdata_offset + $rdlen;
    }

    return @records;
}

sub read_dns_name {
    my ($packet, $offset) = @_;
    my @labels;
    my $pos = $offset;
    my $next;
    my %seen;

    while (1) {
        return unless $pos < length($packet);
        return if $seen{$pos}++;
        my $len = ord(substr($packet, $pos, 1));
        if (($len & 0xc0) == 0xc0) {
            return unless $pos + 1 < length($packet);
            my $ptr = unpack('n', substr($packet, $pos, 2)) & 0x3fff;
            $next = $pos + 2 unless defined $next;
            $pos = $ptr;
            next;
        }
        if ($len == 0) {
            $next = $pos + 1 unless defined $next;
            last;
        }
        return if ($len & 0xc0) || $pos + 1 + $len > length($packet);
        push @labels, substr($packet, $pos + 1, $len);
        $pos += 1 + $len;
    }

    return (join('.', @labels), $next);
}

sub record_is_usable {
    my ($record) = @_;
    return 0 unless ($record->{name} || '') =~ /\A[a-z0-9][a-z0-9_.-]*\.local\z/;
    return 0 unless ip_is_observable($record->{ip} || '');
    return 1;
}

sub ip_is_observable {
    my ($ip) = @_;
    return $ip =~ /\A192\.168\.\d+\.\d+\z/
        || $ip =~ /\A10\.\d+\.\d+\.\d+\z/
        || $ip =~ /\A172\.(?:1[6-9]|2\d|3[01])\.\d+\.\d+\z/
        || $ip =~ /\A169\.254\.\d+\.\d+\z/;
}

sub seed_observations {
    my ($opt, $records, $peer_ip) = @_;
    my $database = $opt->{database};
    ensure_parent_dir($database);
    open my $lock_fh, '+>>', $database or die "Cannot open $database: $!\n";
    flock($lock_fh, LOCK_EX) or die "Cannot lock $database: $!\n";
    seek($lock_fh, 0, 0);
    local $/;
    my $text = <$lock_fh>;
    my $db = parse_observations_yaml($text || '');
    my @changes;

    for my $record (@$records) {
        my $change = merge_observation($db, $record, $peer_ip);
        push @changes, $change if $change;
    }

    if (@changes) {
        if ($opt->{dry_run}) {
            print change_line($_) . "\n" for @changes;
        } else {
            $db->{updated_at} = iso_now();
            seek($lock_fh, 0, 0);
            truncate($lock_fh, 0) or die "Cannot truncate $database: $!\n";
            print {$lock_fh} render_observations_yaml($db);
            close $lock_fh or die "Cannot close $database: $!\n";
        }
    }

    return \@changes;
}

sub merge_observation {
    my ($db, $record, $peer_ip) = @_;
    my $now = iso_now();
    my $id = id_from_mdns_name($record->{name});
    my $key = "$record->{name}|$record->{ip}";
    my $existing;
    for my $observation (@{ $db->{observations} || [] }) {
        if (($observation->{key} || '') eq $key) {
            $existing = $observation;
            last;
        }
    }

    if (!$existing) {
        push @{ $db->{observations} }, {
            key => $key,
            id => $id,
            name => $record->{name},
            ip => $record->{ip},
            first_seen => $now,
            last_seen => $now,
            seen_count => 1,
            last_peer => $peer_ip,
            ttl => int($record->{ttl} || 0),
        };
        return { action => 'created', id => $id, name => $record->{name}, ip => $record->{ip} };
    }

    $existing->{last_seen} = $now;
    $existing->{seen_count} = int($existing->{seen_count} || 0) + 1;
    $existing->{last_peer} = $peer_ip;
    $existing->{ttl} = int($record->{ttl} || 0);
    return { action => 'updated', id => $existing->{id}, name => $record->{name}, ip => $record->{ip} };
}

sub parse_observations_yaml {
    my ($text) = @_;
    my %db = (
        version => 1,
        updated_at => '',
        source => 'mdns',
        observations => [],
    );
    my ($section, $current);
    for my $line (split /\n/, $text) {
        next if $line =~ /^\s*$/ || $line =~ /^\s*#/;
        if ($line =~ /^version:\s*(\d+)/) {
            $db{version} = int($1);
        } elsif ($line =~ /^updated_at:\s*(.+)$/) {
            $db{updated_at} = yaml_unquote($1);
        } elsif ($line =~ /^source:\s*(.+)$/) {
            $db{source} = yaml_unquote($1);
        } elsif ($line =~ /^observations:\s*$/) {
            $section = 'observations';
        } elsif (($section || '') eq 'observations' && $line =~ /^  - key:\s*(.+)$/) {
            $current = { key => yaml_unquote($1) };
            push @{ $db{observations} }, $current;
        } elsif ($current && $line =~ /^    ([A-Za-z0-9_]+):\s*(.*)$/) {
            my ($key, $value) = ($1, yaml_unquote($2));
            $value = int($value || 0) if $key =~ /\A(?:seen_count|ttl)\z/;
            $current->{$key} = $value;
        }
    }
    return \%db;
}

sub render_observations_yaml {
    my ($db) = @_;
    my $out = "version: " . int($db->{version} || 1) . "\n";
    $out .= "updated_at: " . yq($db->{updated_at} || iso_now()) . "\n";
    $out .= "source: " . yq($db->{source} || 'mdns') . "\n";
    $out .= "observations:\n";
    for my $observation (sort { ($a->{name} || '') cmp ($b->{name} || '') || ($a->{ip} || '') cmp ($b->{ip} || '') } @{ $db->{observations} || [] }) {
        $out .= "  - key: " . yq($observation->{key}) . "\n";
        for my $key (qw(id name ip first_seen last_seen seen_count last_peer ttl)) {
            $out .= "    $key: " . yq($observation->{$key} || '') . "\n";
        }
    }
    return $out;
}

sub id_from_mdns_name {
    my ($name) = @_;
    $name =~ s/\.local\z//;
    $name =~ s/\([0-9]+\)\z//;
    $name = lc($name);
    $name =~ s/[^a-z0-9_.-]+/-/g;
    $name =~ s/^-+|-+$//g;
    return $name;
}

sub ensure_parent_dir {
    my ($path) = @_;
    my $dir = dirname($path);
    make_path($dir) unless -d $dir;
}

sub change_line {
    my ($change) = @_;
    return join(' ', map { "$_=$change->{$_}" } sort keys %$change);
}

sub yq {
    my ($value) = @_;
    $value = '' unless defined $value;
    $value =~ s/\\/\\\\/g;
    $value =~ s/"/\\"/g;
    return qq("$value");
}

sub yaml_unquote {
    my ($value) = @_;
    $value = '' unless defined $value;
    $value =~ s/^\s+|\s+$//g;
    if ($value =~ /^"(.*)"$/) {
        $value = $1;
        $value =~ s/\\"/"/g;
        $value =~ s/\\\\/\\/g;
    }
    return $value;
}

sub iso_now {
    return strftime('%Y-%m-%dT%H:%M:%SZ', gmtime);
}