path: root/lib/MGA/Mirrors/DB.pm

package MGA::Mirrors::DB;

# $Id$

use strict;
use warnings;
use Config::IniFiles;
use URI;
use DBI;
use File::Temp qw(tempfile);
use Net::DNS;

sub configfile { '/etc/mga-mirror.ini' }

sub new {
    my ($class) = @_;

    my $conf = (-f './mga-mirror.ini')
        ? Config::IniFiles->new(-file => './mga-mirror.ini')
        : Config::IniFiles->new(-file => configfile())
            or return;

    my $db = DBI->connect(
        'dbi:Pg:' . $conf->val('db', 'pgconn', ''),
        $conf->val('db', 'user') || undef,
        $conf->val('db', 'password') || undef,
        {
            AutoCommit => 0,
            PrintError => 1,
        }
    ) or return;

    bless {
        db => $db,
        conf => $conf,
    }, $class;
}

sub host_ips {
    my ($self, $hostname) = @_;

    my $resolver = Net::DNS::Resolver->new;
    my @addresses;
    foreach my $type (qw'A AAAA') {
        my $packet = $resolver->search($hostname, $type) or next;
        foreach ($packet->answer) {
            $_->type eq $type or next;
            push(@addresses, $_->address);
        }
    }
    @addresses;
}

sub db { $_[0]->{db} }

sub locate_ips {
    my ($self, @ips) = @_;

    my $find = $self->db->prepare(q{
        select countries.* from geoip 
        join countries on geoip.code = countries.code
        where ipmin <= $1 and ipmax >= $1
    });

    foreach (@ips) {
        $find->execute($_);
        my $res = $find->fetchrow_hashref;
        if ($res) {
            $find->finish;
            return $res;
        }
    }

    return;
}

sub country_list {
    my ($self) = @_;
    my $list = $self->db->prepare(q{
        select * from countries order by name
    });
    $list->execute;
    return $list->fetchall_arrayref({});
}

sub mirror_validity {
    my ($self, $uri) = @_;
    my $listf = $self->db->prepare(q{
        select * from global_files
    });
    $listf->execute;
    while (my $res = $listf->fetchrow_hashref) {
        my $furi = URI->new($uri . $res->{relpath});
        $self->_check_url($furi) or return;
    }

    1;
}

sub check_distributions {
    my ($self) = @_;

    my $uneeded_check = $self->db->prepare(q{
        select * from mirrors_distributions where
        lastcheck > now() - '6 hours'::interval
        });
    $uneeded_check->execute();
    my $uch = $uneeded_check->fetchall_hashref([ qw(urlskey distributionkey) ]);

    my $listd = $self->db->prepare(q{
        select * from urls, distributions
        where urls.valid = true
        });

    my $addstatus = $self->db->prepare(q{
        insert into mirrors_distributions (urlskey, distributionkey, exists)
        values (?,?,?)
        });

    my $updstatus = $self->db->prepare(q{
        update mirrors_distributions set lastcheck = now(), exists = ?
        where urlskey = ? and distributionkey = ?
        });

    my %urls_status = ();

    my $updurl = $self->db->prepare(q{
        update urls set lastcheck = now(), valid = ?
        where key = ?
    });

    $listd->execute();
    while (my $res = $listd->fetchrow_hashref) {
        $uch->{$res->{key}}{$res->{dkey}} and next;
        my $url = $self->fmt_url($res);
        if (!exists($urls_status{$res->{key}})) {
            my $ok = $self->mirror_validity($url);
            $updurl->execute($ok ? 1 : 0, $res->{key});
            $urls_status{$res->{key}} = $ok;
        }
        $urls_status{$res->{key}} or next;
        my $furi = URI->new(join('/', $url, $res->{relpath}, $res->{relfile}));
        my $exists = $self->_check_url($furi);
        if ($updstatus->execute($exists, $res->{key}, $res->{dkey}) == 0) {
            $addstatus->execute($res->{key}, $res->{dkey}, $exists);
        }
        $self->db->commit;
    }
}

sub _check_url {
    my ($self, $furi) = @_;
    my ($fh, $filename) = tempfile();
    close($fh);
    my $cmd = 
    $furi->scheme =~ /^http|ftp$/ ? "wget -nv -t 1 -T 4 -O $filename " . $furi->as_string :
    $furi->scheme eq 'rsync' ? "rsync --timeout 4 -q " . $furi->as_string . " $filename" : '';
    my $ok = (system($cmd) == 0);
    unlink($filename);
    return $ok
}

sub get_protocol_info {
    my ($self, $protocol) = @_;
    my $get = $self->db->prepare(q{
        select * from protocol where name = ?
    });
    $get->execute($protocol);
    my $res = $get->fetchrow_hashref;
    $get->finish;
    $res;
}

sub find_mirrors {
    my ($self, $filters, $key) = @_;

    my $query = q{
        select * from hosts
        left join countries on countries.code = hosts.country
        where hosts.hostname in (select hostname from urls %s)
        %s
    };

    my (@mvals, @uvals);
    my (@mw, @uw);
    if (keys %{ $filters || {}}) {
        foreach (keys %$filters) {
            $filters->{$_} or next;
            if (my $field = {
                hostname => 'hosts.hostname',
                country => 'countries.code',
                continent => 'countries.contienent_code',
            }->{$_}) {
                push(@mw, sprintf('%s = any(?)', $field));
                push(@mvals, ref $filters->{$_} ? $filters->{$_} : [ $filters->{$_} ]);
            }
            if (my $field = {
                protocol => 'protocol',
            }->{$_}) {
                push(@uw, sprintf('%s = any(?)', $field));
                push(@uvals, ref $filters->{$_} ? $filters->{$_} : [ $filters->{$_} ]);
            }
        }
    }
    my $list = $self->db->prepare(sprintf(
        $query,
        (@uw ? 'where ' . join(' and ', @uw) : ''),
        (@mw ? 'and '   . join(' and ', @mw) : ''),
    ));
    $list->execute(@uvals, @mvals);
    return $list->fetchall_arrayref({});
}

sub _find_urls {
    my ($self, $filters, $key) = @_;

    my $query = q{
        select urls.* from urls join
        hosts on hosts.hostname = urls.hostname
    };
    my @vals;
    if (keys %{ $filters || {} }) {
        $query .= ' where ';
        my @w;
        foreach (keys %$filters) {
            my $field = {
                hostname => 'hosts.hostname',
                protocol => 'urls.protocol',
            }->{$_} or next;

            push(@w, sprintf('%s = any(?)', $field));
            push(@vals, ref $filters->{$_} ? $filters->{$_} : [ $filters->{$_} ]);
        }
        $query .= join(' and ', @w);
    }
    my $list = $self->db->prepare($query);
    $list->execute(@vals);
    return $list->fetchall_arrayref({});
}

sub find_host_ip_overlap {
    my ($self, $hostname) = @_;

    my @addresses = $self->host_ips($hostname);

    my $list = $self->db->prepare(q{
        select * from ips where ip = any(?)
            and hostname != ?
    });
    $list->execute(\@addresses, $hostname);
    my $res = $list->fetchall_hashref('hostname');
    return keys %{ $res };
}

sub add_or_update_host {
    my ($self, $hostname, %info) = @_;

    my (@fields, @vals);
    while (my ($field, $val) = each(%info)) {
        push(@fields, $field);
        push(@vals,   $val);
    }
    if (keys %info) {
        my $upd = $self->db->prepare(sprintf(q{
            update hosts set %s where hostname = ?
        }, join(', ', map { "$_ = ?" } @fields)));
        if ($upd->execute(@vals, $hostname) == 0) {
            my $add = $self->db->prepare(sprintf(q{
                insert into hosts (%s) values (%s)
                }, join(', ', (@fields, 'hostname')),
                join(',', ('?') x (scalar(@fields)+1))
            ));
            $add->execute(@vals, $hostname) or do {
                $self->db->rollback;
                return;
            };
        }
    }

    $self->update_host_ips($hostname);

    1;
}

sub add_or_update_url {
    my ($self, $uri) = @_;
    if (!ref $uri) {
        $uri = URI->new($uri);
    }

    my $update = $self->db->prepare(q{
        update urls set path = ?, port = ?
        where hostname = ? and protocol = ?
    });

    if ($update->execute(
            $uri->path, $uri->port == $uri->default_port ? undef : $uri->port,
            $uri->host, $uri->scheme
        ) == 0) {
        my $add = $self->db->prepare(q{
            insert into urls (path, port, hostname, protocol)
            values (?,?,?,?)
            });
        $add->execute($uri->path, $uri->port == $uri->default_port ? undef : $uri->port,
            $uri->host, $uri->scheme) or do {
            $self->db->rollback;
            return;
        }
    }

    1;
}

sub update_host_ips {
    my ($self, $hostname) = @_;

    my @addresses = $self->host_ips($hostname);
    my $delete = $self->db->prepare(
        q{delete from ips where hostname = ?
            and ip != any(?)
        }
    );

    $delete->execute($hostname, [ @addresses ]);

    my $getip = $self->db->prepare(q{
        select 1 from ips where hostname = ? and ip = ?
    });
    my $addip = $self->db->prepare(q{
        insert into ips (hostname, ip) values (?,?)
    });
    foreach (@addresses) {
        if ($getip->execute($hostname, $_) == 0) {
            $addip->execute($hostname, $_);
        }
        $getip->finish;
    }
    
    1;
}

sub find_urls {
    my ($self, $filters, $key) = @_;
    return [
        map { $_->{url} = $self->fmt_url($_); $_ }
        @{ $self->_find_urls($filters) || []}]
}

sub fmt_url {
    my ($self, $dburl) = @_;

    my $uri = URI->new(
        sprintf('%s://%s%s',
            $dburl->{protocol},
            $dburl->{hostname},
            $dburl->{path} || '/',
        )        
    );
    $uri->port($dburl->{port});

    return $uri->as_string;
}

1;