#!/usr/bin/env python2

# Zedspider concept
# Takes csv file of ipaddr, port as input
# If service is unique, DB is updated with service info

from gevent import monkey
monkey.patch_all()

import gevent.pool
import gevent.queue

import socket
import html2text
import requests
import sys
import os
import time
import csv
import mysql.connector
import telnetlib
import GeoIP

from randua import randomua

#TODO:
#	Batch update for SQL instead of sequential.

class zedspider_update():

     def __init__(self):
          self.headers={'User-Agent':randomua(),
			'Accept':'text/html, application/xml;q=0.9, application/xhtml+xml, image/png, image/webp, image/jpeg, image/gif, image/x-xbitmap, */*;q=0.1',
			'Connection':'Keep-Alive'}
          self.counter_uip = 0
          self.counter_usvc = 0
          self.counter_esvc = 0

     def run(self):
          # Load the initial list of IPs and ports
          ipdata = self.load_list()

          # Append the geolocation data
          ipdata = self.geo_data(ipdata)
          total_ips = len(ipdata)
          test_queue = []
          # For each row: Build variables, query DB, and update

          while ipdata:
               # Build out variables
               for counter in range(0, min(100, len(ipdata))):
                    qrow = []
                    qrow.append(ipdata[0][0])
                    qrow.append(ipdata[0][1])
                    qrow.append(ipdata[0][2])
                    qrow.append(ipdata[0][3])
                    qrow.append(ipdata[0][4])
                    qrow.append(ipdata[0][5])
                    test_queue.append(qrow)
                    del ipdata[0]

               # SQL batch update code will go here

               jobs = [gevent.spawn(self.process_data, testrow) for testrow in test_queue]
               gevent.joinall(jobs, timeout=30)
               test_queue = []
               print str(self.counter_usvc + self.counter_esvc) + "/" + str(total_ips) + " tested."
               counter = 0

          print "Updated with " + str(self.counter_uip) + " unique IP addresses ; " + str(self.counter_usvc) + " unique services ; " + str(self.counter_esvc) + " existing services."

     def process_data(self, ipdata):
          qrow = []
          qrow.append(ipdata[0])
          qrow.append(ipdata[1])
          qrow.append(ipdata[2])
          qrow.append(ipdata[3])
          qrow.append(ipdata[4])
          qrow.append(time.strftime('%Y/%m/%d %H:%M:%S'))
          qrow.append(self.dump_data(ipdata[0], ipdata[1], 8))
          qrow.append(ipdata[5])
          self.sql_update(qrow[0], qrow[1], qrow[2], qrow[3], qrow[4], qrow[5], qrow[6], qrow[7])

     def load_list(self):
          #read csv into list
          data = []
          filename = sys.argv[1]
          datafile = open(filename, 'r')
          datafile.readline()
          datareader = csv.reader(datafile)
          for row in datareader:
               data.append(row)
          return data

     def geo_data(self, ipdata):
          # Fetch country, region, and city from GeoLite db. Nice and fast.

          gi = GeoIP.open("/usr/share/GeoIP/GeoLiteCity.dat", GeoIP.GEOIP_STANDARD)
          gorg = GeoIP.open("/usr/share/GeoIP/GeoIPOrg.dat", GeoIP.GEOIP_STANDARD)

          counter = 0
          geodata = []
          print "Adding GeoIP information."
          for row in ipdata:
               gir = gi.record_by_addr(row[0])

               geodata.append(row)
               try:
                    geodata[counter].append(str(gir['country_name']))
               except:
                    geodata[counter].append("NA")
               try:
                    geodata[counter].append(str(gir['region_name']))
               except:
                    geodata[counter].append("NA")
               try:
                    geodata[counter].append(str(gir['city']))
               except:
                    geodata[counter].append("NA")
               try:
                    geodata[counter].append(str(gorg.org_by_addr(row[0])))
               except:
                    geodata[counter].append("NA")

               counter += 1

          return geodata

     def geo_data_old(self, ipdata):
          # Fetch country, region, and city from geoiptool.
          # This can also be done with a geolocation db, but
          # for today's purposes this'll work.
          # -- Fuck this. It's slow as fuck. Retained for reference.

          geodata = []
          counter = 0
          cc = 'NA'
          cc_line_found=False
          re = 'NA'
          re_line_found=False
          ct = 'NA'
          ct_line_found=False

          for row in ipdata:
               try:
                self.headers={'User-Agent':randomua()}
               	r = requests.get('http://www.geoiptool.com/en/?IP=%s' % row[0], headers=self.headers)
               	html = r.text
	       except:
		html = "geoip failed"
               html_text = html2text.html2text(html)
               html_lines = html_text.splitlines()

               # Fetch Country Code
               for l in html_lines:
                    if 'country:' in l.lower():
                         try:
                           cc = l.split('gif) ')[1]
                         except:
                           cc = "NA"
                         break

               # Fetch Region/State
               for l in html_lines:
                    if 'region:' in l.lower():
                         try:
                           re = l.split(': ')[1]
                         except:
			   re = "NA"
                         break

               # Fetch City
               for l in html_lines:
                    if 'city:' in l.lower():
                         try:
                           ct = l.split(': ')[1]
                         except:
                           ct = "NA"
                         break

               geodata.append(row)
               geodata[counter].append(cc)
               geodata[counter].append(re)
               geodata[counter].append(ct)

               counter += 1
          return geodata

     def dump_data(self, strIPAddr, intPort, TIMEOUT):
          banner='NA'
          try:
               # Telnet requires different handling. In the future can add multiple protocols.
               if intPort == "23":
                    tn = telnetlib.Telnet(strIPAddr, intPort)
                    ret1 = tn.read_until("Shouldneverseethissojusttimeout", TIMEOUT)
                    banner = ret1
                    tn.close()
               else:
                    self.headers={'User-Agent':randomua(),
                                  'Accept':'text/html, application/xml;q=0.9, application/xhtml+xml, image/png, image/webp, image/jpeg, image/gif, image/x-xbitmap, */*;q=0.1',
                                  'Connection':'Keep-Alive'}

                    r = requests.get('http://'+strIPAddr+':'+str(intPort), headers=self.headers, timeout=TIMEOUT)
                    html = r.text
                    #s=socket.socket()
                    #s.settimeout(TIMEOUT)
                    #s.connect((strIPAddr,int(intPort)))
                    #packet='GET / HTTP/1.0\r\n' + "User-Agent: " + self.headers['User-Agent'] + '\r\nHost: ' + strIPAddr + '\r\n\r\n'
                    #s.send(packet)
                    banner = html
                    #s.close()

               return banner
          except:
               return banner

     def sql_update(self, strIPAddr, intPort, strCC, strRE, strCT, date, banner, strORG):
		cnx = mysql.connector.connect(user='login', password='password', database='database')
		cursor = cnx.cursor()

		res1 = []
		res2 = []

		# Query IP address from IPAddress table and store it in res1
		cursor.execute("select idIPAddr from tbl_IPAddresses where strIPAddr = %s limit 1", (strIPAddr, ))
		res1 = [item[0] for item in cursor.fetchall()]

		# If this IP address doesn't already exist, add its info to the DB
		if not res1:
			#print '%s does not exist, so add it.' % strIPAddr
			self.counter_uip += 1
			add_IP = ("INSERT INTO tbl_IPAddresses (strIPAddr, strCountry, strCity, strRegion, strOrg) "
			          "VALUES (%s, %s, %s, %s, %s)")
			IP_data = (strIPAddr, strCC, strCT, strRE, strORG)
			#print IP_data
			cursor.execute(add_IP, IP_data)
			idIPAddr = cursor.lastrowid

			cnx.commit()
		else:
			idIPAddr = res1[0]


		# Query IP address AND port# from service table and store it in res2
		cursor.execute("select idSvc from tbl_service where (idIPAddr = %s AND intPort = %s) limit 1", (idIPAddr, intPort))
		res2 = [item[0] for item in cursor.fetchall()]

		# If this service (port) doesn't exist, add its info to the DB
		if not res2:
			self.counter_usvc += 1
			#print '%s does not exist for %s, so add it.' % (intPort, strIPAddr)
			add_service = ("INSERT INTO tbl_service (idIPAddr, intPort, dateDiscovered, dateLastChecked) "
				       "VALUES (%s, %s, %s, %s)")
			service_data = (idIPAddr, intPort, date, date)
			#print service_data
			cursor.execute(add_service, service_data)
			idSvc = cursor.lastrowid

			cnx.commit()
		else:
			idSvc = res2[0]
			self.counter_esvc += 1
			#print 'Updating service date last checked.'
			update_service = ("UPDATE tbl_service SET dateLastChecked = %s WHERE (idSvc = %s)")
			update_service_data = (date, idSvc)
			#print update_service_data
			cursor.execute(update_service, update_service_data)

		# Update Service Content
		add_dump = ("INSERT INTO tbl_service_content (idSvc, dateChecked, txtDump) "
			    "VALUES (%s, %s, %s)")
		dump_data = (idSvc, date, banner)
		cursor.execute(add_dump, dump_data)
		cnx.commit()

		# Close cursor
		# Make sure data is committed to the database
		cursor.close()
		cnx.close()

P=zedspider_update()
P.run()