Revert "Revert "windows支持以及linux自动更新""

This reverts commit ad2e4ee.

Author: zyazhb

.github/workflows/CI.yml

@@ -26,7 +26,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install requests fake_useragent ping3 prettytable aiohttp uvloop
+        pip install requests fake_useragent ping3 prettytable aiohttp uvloop crontab
     - name: Build and publish
       run: |
         ./main.py -t github.com

main.py

@@ -2,6 +2,8 @@
 import argparse
 import myutils
 
+import sys
+
 def initArguments():
     parser = argparse.ArgumentParser(
         description="Prevent domains from dns polution.")
@@ -12,7 +14,7 @@ def initArguments():
     parser.add_argument("--area", help="Choose area in china asia europe africa oceania north_america south_america",
                         choices=["china", "asia", "europe", "africa", "oceania", "north_america", "south_america", "debug"], type=str, default="north_america", nargs='?')
     parser.add_argument(
-        "--thread", help="Run num of threads,default is 3", type=int, default=3)
+        "--update", help="Auto update hosts", action="store_true")
     args = parser.parse_args()
     return parser, args
 
@@ -29,6 +31,10 @@ def main():
 
     if args.clean:
         myutils.output_dic(domain, ipdict[1])
+    if args.update and 'win' not in sys.platform:
+        myutils.update_hosts(domain, tuple(ipdict[1].keys()))
+        myutils.update_crontab(domain)
+
 
 if __name__ == "__main__":
     main()

multi.py

@@ -1,12 +1,12 @@
 import asyncio
 import aiohttp
-import uvloop
+
+import sys
 
 import time
 import subprocess
 import re
 
-
 class MyConnector(aiohttp.TCPConnector):
     def __init__(self, ip):
         self.__ip = ip
@@ -32,6 +32,7 @@ def now(): return time.time()
 ip_list = []
 time_list = {}
 
+
 async def run(cmd):
     proc = await asyncio.create_subprocess_shell(
         cmd,
@@ -41,35 +42,57 @@ async def run(cmd):
     stdout, stderr = await proc.communicate()
 
     if stdout:
-        ip_find = re.findall("\\d+\\.\\d+\\.\\d+\\.\\d+", stdout.decode())
+        if platform == "linux":
+            ip_find = re.findall("\\d+\\.\\d+\\.\\d+\\.\\d+", stdout.decode())
+        elif platform == "win":
+            ip_find_tmp = re.findall(
+                "\\d+\\.\\d+\\.\\d+\\.\\d+", stdout.decode("gbk"))
+            real_dns_addr = ip_find_tmp[0]
+            ip_find = [i for i in ip_find_tmp if i != real_dns_addr]
         ip_list.extend(ip_find)
     if stderr:
-        print(f'[stderr]\n{stderr.decode()}')
+        pass
 
 
 async def test_doamin_ip(ip):
     st = now()
     try:
         async with aiohttp.ClientSession(connector=MyConnector(ip), timeout=aiohttp.ClientTimeout(total=10)) as client:
             async with client.get("https://{0}".format(ip), ssl=False, timeout=10) as resp:
-                if resp.status == 200:
+                if resp.status == 200 or 405:
                     time_list[ip] = now() - st
     except asyncio.TimeoutError:
         pass
+    
 
 async def dns_test(domain):
-    task_list = [asyncio.create_task(
-        run('dig @{0} {1} +short'.format(dns, domain))) for dns in dns_list]
+    if platform == "linux":
+        task_list = [asyncio.create_task(
+            run('dig @{0} {1} +short'.format(dns, domain))) for dns in dns_list]
+    elif platform == "win":
+        task_list = [asyncio.create_task(
+            run('nslookup {0} {1}'.format(domain, dns))) for dns in dns_list]
     done, pending = await asyncio.wait(task_list, timeout=5)
 
     task_speed = [asyncio.create_task(test_doamin_ip(ip))
                   for ip in set(ip_list)]
     done, pending = await asyncio.wait(task_speed)
 
-def multi_local_dns(domain):
-    asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
+def multi_local_dns(domain, platform_in):
+    global platform 
+    platform = platform_in
+    if platform == 'linux':
+        import uvloop
+        asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+    elif platform == 'win':
+        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
+    else:
+        raise "platform not support!"
     start = now()
-    asyncio.run(dns_test(domain))
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(dns_test(domain))
+    
+    # asyncio.run(dns_test(domain))
     print("Time: ", now() - start)
     return domain, time_list

myutils.py

@@ -4,14 +4,22 @@
 import re
 from prettytable import PrettyTable
 
+import os
+import sys
+import subprocess
+
+global platform
+if "linux" in sys.platform:
+    from crontab import CronTab
+    platform = "linux"
+elif "win" in sys.platform:
+    platform = "win"
+
 def run_core(domain, area):
     # Encrypt!
-    if area == "debug":
-        iplist = ['220.181.38.148', '39.156.69.79', '210.23.129.34', '210.23.129.34', '220.181.38.148', '39.156.69.79', '202.108.22.220', '220.181.33.31', '112.80.248.64', '14.215.178.80', '180.76.76.92', '210.23.129.34', '210.23.129.34', '39.156.69.79', '220.181.38.148', '203.12.160.35', '203.12.160.35', '39.156.69.79', '220.181.38.148',
-                  '202.108.22.220', '220.181.33.31', '112.80.248.64', '14.215.178.80', '180.76.76.92', '203.12.160.35', '203.12.160.35', '220.181.38.148', '39.156.69.79', '61.8.0.113', '61.8.0.113', '220.181.38.148', '39.156.69.79', '202.108.22.220', '220.181.33.31', '112.80.248.64', '14.215.178.80', '180.76.76.92', '61.8.0.113', '61.8.0.113']
-    else:
-        ipdict = multi.multi_local_dns(domain)
-        print("[+]Got domain! \n" + str(list(ipdict[1].keys())))
+    print("[+]platform detect: "+platform)
+    ipdict = multi.multi_local_dns(domain,platform)
+    print("[+]Got domain! \n" + str(list(ipdict[1].keys())))
     return domain, ipdict
 
 
@@ -47,6 +55,36 @@ def output_dic(domain, ip_dic):
 
     print(table)
 
+def update_hosts(domain, new_ip):
+    if os.getuid() != 0:
+        sys.exit("not root?")
+
+    if len(new_ip) != 0:
+        print("[-]Start updating hosts")
+        for ip in new_ip[::-1]:
+            cmd = ['sed', '-i', rf'/^[0-9.]\+[[:space:]]\+{domain}\>/s/[^[:space:]]\+/{ip}/', '/etc/hosts']
+            try:
+                subprocess.check_call(cmd)
+                print("Add {0} {1}".format(domain, ip))
+            except:
+                print("Error: {0} {1}".format(domain, ip))
+        print("[+]Done!")
+
+def update_crontab(domain):
+    my_user_cron = CronTab(user=True)  # 创建当前用户的crontab
+    # 删除原有的crontab文件中重复的内容
+
+    objs = my_user_cron.find_comment(domain)
+    if objs:
+        for obj in objs:
+            my_user_cron.remove(obj)
+
+    job = my_user_cron.new(command='python3 /program/python/Hosts-chooser-master main.py -t ' + domain + ' --clean')
+    job.setall('*/2 * * * *')  # 设置执行时间
+    job.set_comment(domain)
+
+    my_user_cron.write()
+
 class domainError(Exception):
     def __init__(self,err='invalid domian'):
         Exception.__init__(self,err)

requirements.txt

@@ -1,4 +1,7 @@
+aiohttp==3.6.2
 fake-useragent==0.1.11
 prettytable==0.7.2
 requests==2.23.0
 urllib3==1.25.9
+uvloop==0.14.0
+crontab==0.22.9