#!/usr/bin/env python
# coding: utf-8
# author: Dean Do
import re
import subprocess
import psutil
import time
import redis
import json
import multiprocessing
import datetime
import copy
import xlwt

hosts = 'serverlist'
domain = "https://www.youtube.com"


def sshproxy(action, port, *server):
    if action:
        ip, sshport, user, password = server
        cmd = '''sshpass -p "%s" ssh -qTfnN -D %s %s@%s -p %s''' % (
            password, port, user, ip, sshport)
        p = subprocess.call(cmd, shell=True)
        if not p:
            return True
        else:
            return False
    else:
        for netstat in psutil.net_connections():
            if netstat.laddr[1] == port:
                pid = netstat.pid
                p = psutil.Process(pid)
                p.kill()
                return True
        return


def runtest(ip, sshport, user, password, port, q):
    username = "root"
    sshproxy(0, port)
    proxy = sshproxy('proxy', port, ip, sshport, username, password)
    if proxy:
        proxy = "127.0.0.1:%s" % port
        cmd = "curl -I -o /dev/null -s -w %%{time_total} --socks5-hostname %s %s" % (
            proxy, domain)
        l = []
        for i in range(6):
            p = subprocess.Popen(
                cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            stdout, err = p.communicate()
            if stdout:
                request_time = float(stdout.strip())
                l.append(request_time)
            if err:
                print err.strip()
        raw_data = copy.copy(l)
        l.pop(l.index(max(l)))
        l.pop(l.index(min(l)))
        avg = '%.2f' % (sum(l) / len(l))
        cmd = "curl   --socks5-hostname %s  ipinfo.io" % proxy
        p = subprocess.Popen(
            cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, err = p.communicate()
        if stdout:
            out = eval(stdout.strip())
            country = out.get("country")
            city = out.get("city")
        q.put([country, city, ip, avg, raw_data])
    sshproxy(0, port)


def iplist(hosts):
    l = []
    port = 7070
    with open(hosts, 'r') as f:
        for line in f:
            if re.match('^#.*', line):
                continue
            line = line.strip()
            regex = re.search(
                '^(?P<ip>(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]))\s+(?P<sshport>\d+?)\s+(?P<user>.*?)\s+(?P<sshpass>.*?)$',
                line)
            ip = regex.group('ip')
            sshport = regex.group('sshport')
            password = regex.group('sshpass')
            user = regex.group('user')
            l.append([ip, sshport, user, password, port])
            port += 1
    return l


def xlsprepare():
    titles = ["IP", "COUNTRY", "CITY", "REQUEST TIME TEST 1", "REQUEST TIME TEST 2", "REQUEST TIME TEST 3",
              "REQUEST TIME TEST 4", "REQUEST TIME TEST 5", "REQUEST TIME TEST 6", "REQUEST TIME TEST AVG"]
    workbook = xlwt.Workbook(encoding='ascii')
    worksheet = workbook.add_sheet('Test Result')
    for index, item in enumerate(titles):
        worksheet.write(0, index, label=item)
        width = len(item) + 3
        worksheet.col(index).width = 256 * width
    worksheet.col(0).width = 256 * 15
    worksheet.col(2).width = 256 * 15
    return workbook, worksheet


if __name__ == "__main__":
    tasks = []
    data = []
    q = multiprocessing.Queue()
    for server_info in iplist(hosts):
        ip, sshport, user, password, port = server_info
        tasks.append(multiprocessing.Process(target=runtest,
                                             name=ip, args=(ip, sshport, user, password, port, q)))

    for task in tasks:
        task.start()
        print task.name, "test start..."
        time.sleep(1)
    for task in tasks:
        task.join()

    n = 0
    workbook, worksheet = xlsprepare()
    while True:
        if not q.empty():
            value = q.get(True)
            data.append(value)
            print "\033[1;32;40mIP: %s\033[0m" % value[2]
            print "COUNTRY:", value[0]
            print "CITY:", value[1]
            print "REQUST TIME AVG:", value[3]
            print "=" * 22
            if 1:
                n += 1
                worksheet.write(n, 0, label=value[2])
                worksheet.write(n, 1, label=value[0])
                worksheet.write(n, 2, label=value[1])
                worksheet.write(n, 3, label=value[4][0])
                worksheet.write(n, 4, label=value[4][1])
                worksheet.write(n, 5, label=value[4][2])
                worksheet.write(n, 6, label=value[4][3])
                worksheet.write(n, 7, label=value[4][4])
                worksheet.write(n, 8, label=value[4][5])
                worksheet.write(n, 9, label=float(value[3]))
        else:
            workbook.save('result.xls')
            break