#!/usr/bin/perl -w

# Date: 20150418
# Author: Daniel Marsh
# Email: daniel@stiw.org

#############################################################################
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or (at
## your option) any later version.
## 
## This program is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
## 
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
#############################################################################

# look at batch queries via api. Apparently 4 hashes per second are supported

# define library path - this is for libs installed via CPAN shell on my system
use lib '/home/username/perl5/lib/perl5';

# include the libraries we need
use strict;
use warnings;

use DBI qw(:sql_types); # cannot forget this... big problems otherwise ;)

use VT::API;
use Data::Dumper;
use DBD::SQLite;

# config zone
my $dbfile = "/var/local/ossec-vt.db"; # Location of database to use
my $pathtoschk = "/var/ossec/queue/syscheck/"; # default OSSEC path
my $apikey = "insert private/public api key from VirusTotal";
# end config zone

#################################################################################
# Use the following to create the required SQLite DB
# ---
# BEGIN TRANSACTION;
# CREATE TABLE `files` (
#                     `file_id`       INTEGER PRIMARY KEY AUTOINCREMENT,
#                     `filename`      TEXT,
#                     `checksum`      TEXT
#                    );
# CREATE TABLE `computers` (
#                         `computer_id`   INTEGER PRIMARY KEY AUTOINCREMENT,
#                         `computer`      TEXT,
#                         `file`  TEXT
#                        );
# CREATE TABLE `checksums` (
#                         `checksum_id`   INTEGER PRIMARY KEY AUTOINCREMENT,
#                         `datetime`      TEXT,
#                         `checksum`      INTEGER UNIQUE,
#                         `checked`       INTEGER,
#                         `virus` INTEGER
#                        );
# CREATE TABLE `check_time` (
#                          `last_check`    INTEGER,
#                          `id`    INTEGER
#                         );
# INSERT INTO `check_time` VALUES (NULL,1);
# COMMIT;
# ---
#################################################################################

# open connection to database
my $dbh = DBI->connect("dbi:SQLite:dbname=$dbfile", "", "");

# read the syscheck files into memory - will need to update this to
# loop through files in /var/ossec/queue/syscheck, record the machine
# names as well...

opendir(DIR, $pathtoschk) or die "Could not open $pathtoschk: $!\n";
my @dbs = grep(/\(.*syscheck$/, readdir(DIR));

# hash to store the checksums, files, and machine names
my %checksums; # we're about to populate this...

foreach my $f (@dbs) {
  open(SCHK, $pathtoschk.$f) or die "Could not open file $f: $!\n";

  my $comp = (split(/ /, $f, 2))[0];
  $comp =~ s/\)//;
  $comp =~ s/\(//;

  while (<SCHK>) {
    chomp;                      # remove trailing newline
    # split into three parts
    my @line = split(/ /, $_, 3);
    my @bline = split(/:/, $line[0]);
    # bline[5] is the sha hash, bline[4] is md5
    # line[2] is the filename

    $checksums{$bline[5]}->{$line[2]}->{$comp}=1;
  }
  close(SCHK);
}
close(DIR);

#print Data::Dumper->Dumper(\%checksums);
#print "\n";

# setup new connection to VirusTotal
my $api = VT::API->new(key => $apikey);
my $result;

# get the date and time
my ($sec,$min,$hour,$mday,$mon,$year,$wday,$yday,$isdst) = localtime(time);
$year += 1900; $mon++;
my $now = "$year-$mon-$mday $hour:$min:$sec\n";

#-- # Need to write date/time of last query against VirusTotal into DB,
#-- # then we need to verify that atleast 15 seconds has passed before we
#-- # can make another request
my $sth = "";

my $timesth = $dbh->prepare("UPDATE check_time SET last_check=? WHERE id=1");
#-$timesth->bind_param(1, time(), SQL_INTEGER);
#-$timesth->execute;

my $gettimesth = $dbh->prepare("SELECT last_check FROM check_time WHERE id=1");
$gettimesth->execute;
my $row = $gettimesth->fetch;

foreach my $keys (keys %checksums) {
  my $checked = 0;
  my $virus = 0;
  my $csum = $keys;

  $sth = $dbh->prepare("SELECT checked FROM checksums WHERE checksum = ?");
  $sth->bind_param(1, $csum, SQL_VARCHAR);
  $sth->execute;
  $row = $sth->fetch;

  # we don't want to check any checksums already checked
  if( defined $row->[0] ) {
    print "Skipping $csum\n";
    next if $row->[0] >= 1;
  }

  # get time of last query to VT
  $gettimesth->execute;
  $row = $gettimesth->fetch;

  # check if it's been atleast 15 seconds since last check, or we just
  # goto sleep for 20 seconds
  my $t = time();
  my $diff = $t - $row->[0];
  if ( $diff <= 15 ) {
    print "Sleeping for 20...\n";
    sleep(20);
  }

  # check VT for the reported CSUM
  $result = $api->get_file_report($csum);
  print "Checked $csum at ".time()."\n";
  if(not defined $result->{"report"} ) {
    $checked = 2;
    # set it to 2 to indicate that no report was obtained
  }
  # set time of this check
  $timesth->bind_param(1, time(), SQL_INTEGER);
  $timesth->execute;

  # check for viruses/malware
  my $hashref = $result->{"report"}->[1];
  slee(p16); # Can only do 4 queries a minute...

  # set the virus counter for the number of items that are detected as viruses
  my $virus_counter = 0;
  foreach my $key ( keys %{$hashref} ) {
    if ( $hashref->{$key} ne '' ) {
      $virus_counter++;
    }
    # if we got here, we've checked that csum
    $checked=1;
  }

  # report if virus detected
  $virus = $virus_counter > 0 ? 1 : 0;

  # checksum column has UNIQUE constraint...
  # so make sure all found checksums are inserted into the database
  $sth = $dbh->prepare("INSERT OR REPLACE INTO checksums ('checksum_id', 'datetime', 'checksum', 'checked','virus') VALUES (NULL, ?, ?, ?, ?)");
  $sth->bind_param(1, $now, SQL_VARCHAR);
  $sth->bind_param(2, $csum, SQL_VARCHAR);
  $sth->bind_param(3, $checked, SQL_INTEGER);
  $sth->bind_param(4, $virus, SQL_INTEGER);
  $sth->execute;

  # populate the files table
  foreach my $files ( keys %{$checksums{$keys}}) {
    $sth = $dbh->prepare("DELETE FROM files WHERE filename=? AND checksum=?");
    $sth->bind_param(1, $files, SQL_VARCHAR);
    $sth->bind_param(2, $csum, SQL_VARCHAR);
    $sth->execute;
    $sth = $dbh->prepare("INSERT OR IGNORE INTO files ('file_id', 'filename', 'checksum') VALUES (NULL, ?, ?)");
    $sth->bind_param(1, $files, SQL_VARCHAR);
    $sth->bind_param(2, $csum, SQL_VARCHAR);
    $sth->execute;

    # populate the computers table
    foreach my $comp (keys %{$checksums{$keys}->{$files}}) {
      $sth = $dbh->prepare("DELETE FROM computers WHERE computer=? and file=?");
      $sth->bind_param(1, $comp, SQL_VARCHAR);
      $sth->bind_param(2, $files, SQL_VARCHAR);
      $sth->execute;
      $sth = $dbh->prepare("INSERT INTO computers ('computer_id', 'computer', 'file') VALUES (NULL, ?, ?)");
      $sth->bind_param(1, $comp, SQL_VARCHAR);
      $sth->bind_param(2, $files, SQL_VARCHAR);
      $sth->execute;
    }
  }
}

#-- # eicar signature for testing
#-- my $eicar = "131f95c51cc819465fa1797f6ccacf9d494aaaff46fa3eac73ae63ffbdfd8267";
#-- # clean signature for testing
#-- my $clean = "30ced927e23b584725cf16351394175a6d2a9577";

exit;