#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/9/3 23:45
# @Author  : leyan
# @File    : zoomeye_v2.0.py
# @Software: PyCharm

import os
import json
import time
import math
import openpyxl
import requests

#获取zoomeye认证token
def zoomeye_login(username,password):
    url = 'https://api.zoomeye.org/user/login'
    jsondata = {'username':username,'password':password}
    json_enconde = json.dumps(jsondata)
    respond = requests.post(url,data=json_enconde)
    try:
        json_access = json.loads(respond.text)
        access_token = json_access['access_token']
    except :
        return False
    return access_token


def zoomeye_write_xls(ip,country_code,city_name,site,title):
    if  not os.path.exists("信息表.xlsx"):
        title_list = ['ip','title','url']
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = '信息汇总'
        for column_num in range(1,len(title_list)+1):
            ws.cell(row=1,column=column_num,value=title_list[column_num-1])
        wb.save('信息表.xlsx')
        print('[-]信息表创建成功。')
    if country_code !=  'CN' or (country_code == 'CN' and city_name =='香港'):
        data_list =['  '.join(ip),title,site]
        wb = openpyxl.load_workbook('信息表.xlsx')
        ws = wb.active
        max_row = ws.max_row       
        for c in range(1,len(data_list)+1):
            ws.cell(row=max_row+1,column=c,value=data_list[c-1])
        wb.save('信息表.xlsx')

def zoomeye_write(matches):
    for m in matches:
        ip = m['ip']
        try:
            country_code = m['geoinfo']['country']['code']
        except:
            country_code = 'null'
        try:
            city_name = m['geoinfo']['city']['names']['zh-CN']
        except:
            city_name = 'null'
        site = m['site']
        title = m['title']
        zoomeye_write_xls(ip,country_code,city_name,site,title)

#返回搜索的总页数
def zoomeye_page_num(username,password,title):
    url = 'https://api.zoomeye.org/web/search?query=title:'+title+'&page=1'
    if not zoomeye_login(username, password):
        print("[-]请输入正确的账号或密码")
        return -1
    headers = {'Authorization': 'JWT ' + zoomeye_login(username, password)}
    time.sleep(2)
    respond_obj = requests.get(url,headers=headers)
    matchs = json.loads(respond_obj.text)['matches']
    if not matchs:
        with open('./log.txt', 'a', encoding='utf8') as f:
            f.write("[-]" + url +"资源用完" + "\n")
        return -2
    total = json.loads(respond_obj.text)['total']
    zoomeye_write(matchs)
    if total == 0:
        return 0
    return math.ceil(total/len(matchs))

#关键词搜索
def zoomeye_search(username,password,title,page_start_num = 1):
    headers = {'Authorization': 'JWT ' + zoomeye_login(username,password)}
    page_num = zoomeye_page_num(username, password, title)
    if page_num == -1:
        return
    elif page_num == 0:
        print('[-]搜索结果为0')
    elif page_num == -2:
        print('[-]资源用完')
    for num in range(page_start_num,page_num+1):
        time.sleep(2)
        url = 'https://api.zoomeye.org/web/search?query=title:'+title+'&page='+ str(num)
        respond_html = requests.get(url,headers=headers)
        try:
            matchs = json.loads(respond_html.text)['matches']
        except:
            print('[-]provided credentials do not provide access to specified resource')
            return
        zoomeye_write(matchs)
        print("[-]",url,"完成下载")
        with open('./log.txt', 'a', encoding='utf8') as f:
            f.write("[-]"+url+"完成下载" + "\n")

if __name__ == '__main__':
    print("###############################")
    print("     @Author : leyan")
    print("    @日期 ： 2019年9月1")
    print("  @服务支持：1075180052(QQ)")
    print(" @仅供学习交流使用，请勿传播！")
    print("###############################")
    user = input('[-] PLEASE INPUT YOUR USERNAME:')
    passwd = input('[-] PLEASE INPUT YOUR PASSWORD:')
    tittle = input('[-] PLEASE INPUT SEARCH KEYWORDS:')
    page_num = input('[-] PLEASE INPUT PAGE START NUM:(DEFAULT=1)')
    if not page_num:
        page_num = 1
    zoomeye_search(user,passwd,tittle,int(page_num))
    os.system('pause')