################################################################################
# Name: pushtodetect.py
#
################################################################################

import ssl
import sys
import json
import base64
import urllib
import urllib2
import logging
import argparse


def groups_2_nameid(current_groups):
   name_id = {}
   for group_element in current_groups:
      name_id[group_element["Name"]] = group_element["Id"]
   return name_id


def auth(username,password):
   return base64.encodestring('%s:%s' % (username, password)\
                              ).replace('\n', '')


def get_session(base_url,username,password):
   logging.debug("get_session()")
   logging.debug("-----------------")
   req = urllib2.Request('%s/auth' % base_url)
   req.add_header("Authorization", "Basic %s" % auth(username,password))
   res = urllib2.urlopen(req)
   logging.debug(res.getcode())
   logging.debug("-----------------")
   return res.read()


def get_groups(session,api_path):
   logging.debug("get_groups()")
   logging.debug("-----------------")
   req = urllib2.Request('%s/groups' % api_path)
   req.add_header('session', session)
   req.add_header('tanium-session', session)
   res = urllib2.urlopen(req)
   jread = json.loads(res.read())
   logging.debug( res.getcode())
   logging.debug("-----------------")
   logging.debug(json.dumps(jread, indent=4))
   logging.debug("-----------------")
   return jread


def add_group(session, group_name, api_path):
   logging.debug("add_group()")
   logging.debug("-----------------")
   req = urllib2.Request('%s/groups' % api_path)
   req.add_header('session', session)
   req.add_header('tanium-session', session)
   data = urllib.urlencode({
      'Name': group_name,
      'Id': len(get_groups(session,api_path)) +1
   })
   res = urllib2.urlopen(req, data=data)
   logging.debug( res.getcode() )
   logging.debug("-----------------")
   return res


def add_rule(session, api_path, new_ioc, group_id):
   logging.debug("add_rule()")
   logging.debug("-----------------")
   
   fields = {
         "GroupId": group_id
         }
      
   rule = {
         'filename': 'generated.rule',
         'content': new_ioc
        }
      
   delimiter = '----------lImIt_oF_tHe_FiLe_12_$'
   crlf = '\r\n'
   
   # Build multipart body
   body = [
      '--' + delimiter,
      'Content-Disposition: form-data; name="0"',
      '',
      json.dumps(fields),
      '--' + delimiter,
      'Content-Disposition: form-data; name="1"; filename="%s"'\
                                            % rule['filename'],
      'Content-Type: application/octet-stream',
      '',
      rule['content'],
      '--' + delimiter + '--'
      ]
   
   data = crlf.join(body)
   
   # You need multipart content-type header
   headers = {
      'session': session,
      'tanium-session': session,
      'content-type': 'multipart/form-data; boundary=%s' % delimiter
   }

   req = urllib2.Request('%s/iocs' % api_path, data=data, headers=headers)
   res = urllib2.urlopen(req)
   
   logging.debug(res.getcode())
   logging.debug("-----------------")
   return res

################################################################################
# Main
################################################################################
if __name__ == "__main__":
   
#parse our arguments
   parser = argparse.ArgumentParser(description='Feed an ioc to Detect')
   
   parser.add_argument('--tanium_username',
                       required = True,
                       default   = None,
                       help     = 'Tanium username')
   
   parser.add_argument('--tanium_password',
                       required = True,
                       default  = None,
                       help     = 'Tanium password')
   
   parser.add_argument('--tanium_host',
                       required = True,
                       default  = None,
                       help     = 'Tanium hostname')
   
   parser.add_argument('--tanium_port',
                       required = False,
                       default  = "443",
                       help     = 'Alternative tanium host port')
   
   parser.add_argument('--ioc_file',
                       required = False,
                       default  = None,
                       help     = 'The ioc file to upload to Tanium')
   
   parser.add_argument('--group_name',
                       required = True,
                       default  = None,
                       help     = 'Group to upload to Detect')
   
   parser.add_argument('--debug',
                       action   ='store_true',
                       required = False,
                       default  = None,
                       help     = 'Enable debug log')
   
   #move the args to vars as an arg dict
   args = vars(parser.parse_args())
   
   #if requested switch on debug logs
   if args['debug'] == True:
      logging.basicConfig(filename='pushtodetect.log',level=logging.DEBUG)
   
   #Ignore the ssl cert when we get one
   if hasattr(ssl, '_create_unverified_context'):
      ssl._create_default_https_context = ssl._create_unverified_context
   
   #set up our base an path
   base_url = 'https://%s:%s' % (args['tanium_host'],args['tanium_port'])
   api_path = '%s/plugin/products/detect/api/v2' % base_url
   
   #The ioc should be in a file or stdin. Slurp it up
   if args['ioc_file'] != None:
      with open(args['ioc_file']) as x: new_ioc = x.read()
   else:
      new_ioc = sys.stdin.read()
   
   #Auth to the service
   session = get_session(base_url,args['tanium_username'],\
                         args['tanium_password'])
   
   #Get a dict of the current groups. 
   current_groups = get_groups(session,api_path)
   
   #convert our json to 'name'='id' dict
   name_id = groups_2_nameid(current_groups)
   
   #check to see if we have or name in this dict
   if  args['group_name'] not in name_id.keys():
      #if we are missing our name then add it
      add_group(session, args['group_name'], api_path)
      current_groups = get_groups(session,api_path)
      name_id = groups_2_nameid(current_groups)
   
   #Upload the ioc to the group id number
   group_id = name_id[args['group_name']]
   add_rule(session, api_path, new_ioc, group_id)