#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2011-2015 Cesnet z.s.p.o
# Use of this source is governed by a 3-clause BSD-style license, see LICENSE file.

from warden_client import Client, Error, read_cfg
import json
import string
import sys
from time import time, gmtime, strftime, sleep
from math import trunc
from uuid import uuid4
from os import path
import base64
import pprint

import sqlite3

DEFAULT_ACONFIG = 'warden_client-dio.cfg'
DEFAULT_WCONFIG = 'warden_client.cfg'
DEFAULT_BINPATH = '/opt/dionaea/var/dionaea/binaries'
DEFAULT_DBFILE  = '/opt/dionea/var/dionea/logsql.sqlite'
DEFAULT_NAME = 'org.example.warden.test'
DEFAULT_AWIN = 5
DEFAULT_CON_ATTEMPTS = 3
DEFAULT_CON_RETRY_INTERVAL = 5
DEFAULT_ATTACH_NAME = 'att1'


def get_precise_timestamp(epoch=None):
    t = epoch if epoch else time()
    us = trunc((t-trunc(t))*1000000)
    g = gmtime(t)
    iso = '%04d-%02d-%02dT%02d:%02d:%02d.%0dZ' % (g[0:6]+(us,))
    return iso

def gen_attach_idea(logger, binaries_path, filename, hashtype, hashdigest, vtpermalink, avref):
    
  refs = []
  if vtpermalink is not None:
    refs.append('url:' + vtpermalink)
  
  if avref is not None:
    refs.extend(avref.split(';'))
  
  if refs:
    refs = list(set(refs))

  try:
    fpath = path.join(binaries_path, hashdigest)
    with open(fpath, "r") as f:
      fdata = f.read()
  except (IOError) as e:
    logger.info("Reading id file \"%s\" with malware failed, information will not be attached." % (fpath))
    return None
    
  attach = { 
    "Attach": [
      {
         "Handle": DEFAULT_ATTACH_NAME,
         "FileName": [filename],
         "Type": ["Malware"],
         "ContentType": "application/octet-stream",
         "Hash": ["%s:%s" % (hashtype, hashdigest)],
         "Size": len(fdata),
         "Ref": refs,
         "ContentEncoding": "base64",
         "Content": base64.b64encode(fdata)
      }
   ]
  }

  return attach

def gen_event_idea(logger, binaries_path, client_name, detect_time, win_start_time, win_end_time, aggr_win, **data):

  category = []
  event = {
     "Format": "IDEA0",
     "ID": str(uuid4()),
     "DetectTime": detect_time,
     "WinStartTime": win_start_time,
     "WinEndTime": win_end_time,
     "ConnCount": data['attack_scale'],
     "Source": [{}],
     "Target": [{}],
     "Node": [
        {
           "Name": client_name,
           "Tags": ["Connection","Honeypot","Recon"],
           "SW": ["Dionaea"],
           "AggrWin": strftime("%H:%M:%S", gmtime(aggr_win * 60))
        }
     ]
  }

  # Determine IP address family
  af = "IP4" if not ':' in data['src_ip'] else "IP6"
  
  # event['Source'].append( af : [data['src_ip']] )
  event['Source'][0][af] = [data['src_ip']]
  event['Target'][0][af] = [data['dst_ip']]

  # Extract & save proto and service name
  proto = [data['proto']]

  if data['service'] in ['mysql', 'mssql']:
    proto.append(data['service'])
  elif data['service'] in ['httpd', 'smbd']:
    proto.append(data['service'][:-1])

  event['Target'][0]['Proto'] = proto

  # Choose correct category
  if data['service'] != 'pcap':
    category.append('Attempt.Exploit')
  else:
    category.append('Recon.Scanning')

  # SMB allows save malware
  if data['service'] == 'smbd' and data['download_md5_hash'] is not None:
    event['Source'][0]['URL'] = [data['download_url']]

    filename = data['download_url'].split('/')[-1]
    
    # Generate "Attach" part of IDEA
    a = gen_attach_idea(logger, binaries_path, filename, "md5", data['download_md5_hash'], data['virustotal_permalink'], data['scan_result'])
    
    # Attach malware section if file with was loaded successfully
    if a is not None:
      category.append('Malware')
      event['Source'][0]['AttachHand'] = [DEFAULT_ATTACH_NAME]
      event.update(a)

  event['Source'][0]['Port'] = [data['src_port']]
  event['Target'][0]['Port'] = [data['dst_port']]
  event['Category'] = category

  pprint.pprint(event)
  return event

def main():
    aconfig = read_cfg(DEFAULT_ACONFIG)
    wconfig = read_cfg(aconfig.get('warden', DEFAULT_WCONFIG))
    
    aname = aconfig.get('name', DEFAULT_NAME)
    awin = aconfig.get('awin', DEFAULT_AWIN)
 
    abinpath = aconfig.get('binaries_path', DEFAULT_BINPATH)
    adbfile = aconfig.get('dbfile', DEFAULT_DBFILE)
    aconattempts = aconfig.get('con_attempts', DEFAULT_CON_ATTEMPTS)
    aretryinterval = aconfig.get('con_retry_interval', DEFAULT_CON_RETRY_INTERVAL)
    
    wconfig['name'] = aname

    wclient = Client(**wconfig)   

    con = sqlite3.connect(adbfile)
    con.row_factory = sqlite3.Row
    crs = con.cursor()

    events = []
    
    # query =  "SELECT c.connection_timestamp AS timestamp, c.remote_host AS src_ip, c.remote_port AS src_port, c.connection_transport AS proto, \
    #           c.local_host AS dst_ip, c.local_port AS dst_port, COUNT(c.connection) as attack_scale, c.connection_protocol AS service, d.download_url, d.download_md5_hash, \
    #           v.virustotal_permalink, GROUP_CONCAT('urn:' || vt.virustotalscan_scanner || ':' || vt.virustotalscan_result,';') AS scan_result \
    #           FROM connections AS c LEFT JOIN downloads AS d ON c.connection = d.connection \
    #           LEFT JOIN virustotals AS v ON d.download_md5_hash = v.virustotal_md5_hash \
    #           LEFT JOIN virustotalscans vt ON v.virustotal = vt.virustotal \
    #           WHERE datetime(connection_timestamp,'unixepoch') > datetime('now','-%d minutes') AND c.remote_host != '' \
    #           GROUP BY c.remote_host, c.local_port ORDER BY c.connection_timestamp ASC;" % (awin)

    query =  "SELECT c.connection_timestamp AS timestamp, c.remote_host AS src_ip, c.remote_port AS src_port, c.connection_transport AS proto, \
              c.local_host AS dst_ip, c.local_port AS dst_port, COUNT(c.connection) as attack_scale, c.connection_protocol AS service, d.download_url, d.download_md5_hash, \
              v.virustotal_permalink, GROUP_CONCAT('urn:' || vt.virustotalscan_scanner || ':' || vt.virustotalscan_result,';') AS scan_result \
              FROM connections AS c LEFT JOIN downloads AS d ON c.connection = d.connection \
              LEFT JOIN virustotals AS v ON d.download_md5_hash = v.virustotal_md5_hash \
              LEFT JOIN virustotalscans vt ON v.virustotal = vt.virustotal \
              WHERE c.remote_host != '' \
              GROUP BY c.remote_host, c.local_port ORDER BY c.connection_timestamp ASC;"

    attempts = 0
    while attempts < aconattempts:
      try:
        crs.execute(query)
        break
      except sqlite3.Error, e:
        attempts += 1
        wclient.logger.info("Info: %s - attempt %d/%d." % (e.args[0], attempts, aconattempts))
        if attempts == aconattempts:
          wclient.logger.error("Error: %s (dbfile: %s)" % (e.args[0], adbfile))

        sleep(aretryinterval)

    rows = crs.fetchall()

    if con:
      con.close

    etime = get_precise_timestamp(time())
    stime = get_precise_timestamp(time() - awin * 60)

    for row in rows:
      dtime = get_precise_timestamp(row['timestamp'])
      events.append(gen_event_idea(logger = wclient.logger, binaries_path = abinpath, client_name = aname, detect_time = dtime, win_start_time = stime, win_end_time = etime, aggr_win = awin, **row))
        
    # print "=== Sending ==="
    # start = time()
    # ret = wclient.sendEvents(events)
    
    # if ret:
    #   wclient.logger.info("%d event(s) successfully delivered." % len(rows))

    # print "Time: %f" % (time() - start)
  
  

if __name__ == "__main__":
    main()