#!/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);

}