#!/usr/bin/perl -w

use strict;
use warnings;
use JSON;
use Getopt::Long;
use LWP::UserAgent;
use HTTP::Cookies;
use URI;
use Data::Dumper;

umask 077;

my $user      = 'zabbix';
my $pass      = 'secret';
my $site      = 'default';
my $url       = 'https://localhost:8443';
my $certcheck = 1;
my $unifi;
my $dev;
my $station;
my $net;
my $wlan;
my $pretty    = 0;

my $json = {};
my $site_id;

GetOptions (
  'user=s'       => \$user,
  'password|p=s' => \$pass,
  'site=s'       => \$site,
  'url=s'        => \$url,
  'cert-check!'  => \$certcheck,
  'unifi'        => \$unifi,
  'dev=s'        => \$dev,
  'station=s'    => \$station,
  'net=s'        => \$net,
  'wlan=s'       => \$wlan,
  'pretty'       => \$pretty
);

# If connecting to localhost, no need to check certificate
my $uri = URI->new($url);
if ($uri->host =~ m/^localhost|127\.0\.0/){
  $certcheck = 0;
}

my @radio_proto = qw/a b g na ng ac/;
my $resp;
my $username = $ENV{LOGNAME} || $ENV{USER} || getpwuid($<);
my $cj = HTTP::Cookies->new(
  file           => "/tmp/.unifi_$username.txt",
  autosave       => 1,
  ignore_discard => 1
);

my $sslopts = {};
if (not $certcheck){
  $sslopts = { verify_hostname => 0, SSL_verify_mode => 0 }
}
my $ua = LWP::UserAgent->new(
  ssl_opts => $sslopts,
  cookie_jar => $cj
);

# Check if we need to login
$resp = $ua->get($url . '/api/self/sites');
if ($resp->is_error){
  # Log into the API
  $resp = $ua->post(
    $url . '/api/login',
    Content => to_json({ username => $user, password => $pass }),
    Content_Type => 'application/json;charset=UTF-8'
  );
  die "Login failed: " . $resp->message . "\n" if $resp->is_error;
  $resp = $ua->get($url . '/api/self/sites');
  die $resp->message . "\n" if $resp->is_error;
}

# Now, we need to get the site ID
foreach (@{from_json($resp->decoded_content)->{data}}){
  if ($_->{name} eq $site || $_->{desc} eq $site){
    $site_id = $_->{_id};
    # If site is referenced by description, translate it to name
    $site = $_->{name} if ($_->{name} ne $site);
    last;
  }
}
die "Site $site not found\n" unless ($site_id);

# Global info about the instance of Unifi
if ($unifi){
  $resp = $ua->get($url . '/api/s/' . $site . '/stat/health');
  die "ZBX_NOTSUPPORTED\n" if $resp->is_error;
  foreach my $entry (@{from_json($resp->decoded_content)->{data}}){
    if ($entry->{subsystem} eq 'wlan'){
      $json->{wireless_clients} = $entry->{num_user};
      $json->{wireless_guests}  = $entry->{num_guest};
    } elsif ($entry->{subsystem} eq 'lan'){
      $json->{wired_clients} = $entry->{num_user};
      $json->{wired_guests}  = $entry->{num_guest};
    }
    foreach (qw/adopted pending disabled/){
      $json->{'dev_' . $_} += $entry->{'num_' . $_} if (defined $entry->{'num_' . $_});
    }
    foreach (qw/num_ap num_sw num_gw/){
      $json->{$_} += $entry->{$_} if ($entry->{$_});
    }
  }
  $json->{$_} ||= 0 foreach (qw/wireless_clients wireless_guests
                                wired_clients wired_guests dev_adopted
                                dev_pending dev_disabled num_ap num_sw
                                num_gw/);
  $resp = $ua->get($url . '/api/s/' . $site . '/stat/sysinfo');
  die "ZBX_NOTSUPPORTED\n" if $resp->is_error;
  $json->{$_} = from_json($resp->decoded_content)->{data}->[0]->{$_}
    foreach (qw/version build update_available/);

  # Get unarchived alarms
  $resp = $ua->post($url . '/api/s/' . $site . '/stat/alarm',
    Content => to_json({ archived => 'false' }),
    Content_Type => 'application/json;charset=UTF-8'
  );
  die "ZBX_NOTSUPPORTED\n" if $resp->is_error;
  $json->{alarm} = scalar @{from_json($resp->decoded_content)->{data}};

} elsif ($dev) {
  # Dev is identified by MAC
  $resp = $ua->get($url . '/api/s/' . $site . '/stat/device/' . $dev);
  die "ZBX_NOTSUPPORTED\n" if $resp->is_error;
  my $obj = from_json($resp->decoded_content)->{data}->[0];
  foreach (qw/sys_stats locating serial name num_sta user-num_sta
              guest-num_sta inform_url version model state type
              cfgversion adopted avg_client_signal/){
    $json->{$_} = $obj->{$_} if (defined $obj->{$_});
  }
  # Convert last seen into a relative time
  $json->{last_seen} = (defined $obj->{last_seen}) ? time - $obj->{last_seen} : time;
  # Add some more info in sys_stats
  $json->{sys_stats}->{$_} = $obj->{'system-stats'}->{$_} foreach (qw/cpu mem uptime/);

  # If this is an ap
  if ($obj->{type} eq 'uap'){

    foreach (qw/guest-rx_packets guest-tx_packets guest-rx_bytes
                guest-tx_bytes user-rx_packets user-tx_packets
                user-rx_bytes user-tx_bytes rx_packets tx_packets
                rx_bytes tx_bytes rx_errors tx_errors
                rx_dropped tx_dropped/){
      $json->{net_stats}->{$_} = $obj->{stat}->{ap}->{$_} if (defined $obj->{stat}->{ap}->{$_});
    }

    # Count the number of SSID served
    $json->{num_wlan} = scalar @{$obj->{radio_table}};

    $resp = $ua->get($url . '/api/s/' . $site . '/stat/sta');
    die "ZBX_NOTSUPPORTED\n" if $resp->is_error;

    foreach my $proto (@radio_proto){
      $json->{$_ . $proto} = 0 foreach (qw/num_sta_ avg_rx_rate_ avg_tx_rate_/);
    }

    foreach my $entry (@{from_json($resp->decoded_content)->{data}}){
      next if (not $entry->{ap_mac} or $entry->{ap_mac} ne $dev or $entry->{is_wired} == JSON::true);
      foreach (@radio_proto){
        if ($entry->{radio_proto} eq $_){
          $json->{'num_sta_' . $_}++;
          $json->{'avg_rx_rate_' . $_} += $entry->{rx_rate};
          $json->{'avg_tx_rate_' . $_} += $entry->{tx_rate};
        }
      }
      $json->{$_} += $entry->{$_} foreach (qw/rx_bytes tx_bytes rx_packets tx_packets/);
      $json->{'avg_' . $_} += $entry->{$_} foreach (qw/satisfaction tx_power signal noise/);
    }
  
    # Now lets compute average values
    $json->{'avg_' . $_} = ($json->{num_sta} == 0) ? undef : $json->{'avg_' . $_} / $json->{num_sta}
                                                  foreach (qw/satisfaction tx_power signal noise/);
    foreach my $proto (@radio_proto){
      $json->{'avg_' . $_ . '_rate_' . $proto} = ($json->{'num_sta_' . $proto} == 0) ?
        undef : $json->{'avg_' . $_ . '_rate_' . $proto} / $json->{'num_sta_' . $proto}
                                                foreach (qw/tx rx/);
    }
  } elsif ($obj->{type} eq 'usw'){
    foreach (qw/rx_packets tx_packets
                rx_bytes tx_bytes rx_errors tx_errors
                rx_dropped tx_dropped/){
      $json->{net_stats}->{$_} = $obj->{stat}->{sw}->{$_} if (defined $obj->{stat}->{sw}->{$_});
    }
  }
} elsif ($station) {
  # Client is identified by MAC
  $resp = $ua->get($url . '/api/s/' . $site . '/stat/sta/' . $station);
  die "ZBX_NOTSUPPORTED\n" if $resp->is_error;
  my $obj = from_json($resp->decoded_content)->{data}->[0];
  my @client_base = qw/rx_packets tx_packets rx_bytes tx_bytes hostname last_seen ip authorized oui is_guest/;
  foreach (@client_base){
    $json->{$_} = $obj->{$_} || 0;
  }
  # Convert last_seen to relative
  $json->{last_seen} = (defined $obj->{last_seen}) ? time - $obj->{last_seen} : time;

  # For wireless stations, we gather some more info
  if ($obj->{is_wired} == JSON::false){
    my @client_wireless = qw/rx_rate tx_rate essid ap_mac tx_power radio_proto signal noise satisfaction/;
    foreach (@client_wireless){
      $json->{$_} = $obj->{$_} || 0;
    }
    # We have the MAC of the AP, lets try to find the name of this AP
    $resp = $ua->get($url . '/api/s/' . $site . '/stat/device/' . $json->{ap_mac});
    die "ZBX_NOTSUPPORTED\n" if $resp->is_error;
    $json->{ap} = from_json($resp->decoded_content)->{data}->[0]->{name};
  }

} elsif ($wlan) {
  # Wlan is identified by ID
  $resp = $ua->get($url . '/api/s/' . $site . '/rest/wlanconf/' . $wlan);
  die "ZBX_NOTSUPPORTED\n" if $resp->is_error;
  my $obj = from_json($resp->decoded_content)->{data}->[0];
  foreach (qw/name security mac_filter_policy vlan/){
    $json->{$_} = $obj->{$_};
  }
  # For boolean, we need to convert
  foreach (qw/enabled is_guest mac_filter_enabled vlan_enabled/){
    $json->{$_} = (defined $obj->{$_} and $obj->{$_} == JSON::PP::true) ? 1 : 0;
  }

  # Now, we need to count stations for each SSID
  $resp = $ua->get($url . '/api/s/' . $site . '/stat/sta');
  die "ZBX_NOTSUPPORTED\n" if $resp->is_error;

  # Set default values to 0
  $json->{num_sta} = 0;
  $json->{'num_sta_' . $_} = 0 foreach (@radio_proto);
  $json->{$_} = 0 foreach (qw/rx_bytes tx_bytes rx_packets tx_packets/);

  foreach my $entry (@{from_json($resp->decoded_content)->{data}}){
    next if (not $entry->{essid} or $entry->{essid} ne $json->{name} or $entry->{is_wired} == JSON::PP::true);
    $json->{num_sta}++;
    foreach (@radio_proto){
      if ($entry->{radio_proto} eq $_){
        $json->{'num_sta_' . $_}++;
        $json->{'avg_rx_rate_' . $_} += $entry->{rx_rate};
        $json->{'avg_tx_rate_' . $_} += $entry->{tx_rate};
      }
    }
    $json->{$_} += $entry->{$_} foreach (qw/rx_bytes tx_bytes rx_packets tx_packets/);
    $json->{'avg_' . $_} += $entry->{$_} foreach (qw/satisfaction tx_power signal noise/);
  }

  # Now lets compute average values
  $json->{'avg_' . $_} = ($json->{num_sta} == 0) ? undef : $json->{'avg_' . $_} / $json->{num_sta}
                                                foreach (qw/satisfaction tx_power signal noise/);
  foreach my $proto (@radio_proto){
    $json->{'avg_' . $_ . '_rate_' . $proto} = ($json->{'num_sta_' . $proto} == 0) ?
      undef : $json->{'avg_' . $_ . '_rate_' . $proto} / $json->{'num_sta_' . $proto}
                                                foreach (qw/tx rx/);
  }
}

print to_json($json, { pretty => $pretty });