#!/usr/bin/evn python
# -*- coding:utf-8 -*-

import os
import time
import sys
import subprocess
import re
import requests

sys.path.append(r'./bin/')
script_dir = os.path.dirname(__file__)
rel_path = r'bin\tshark'
tshark_file_path = os.path.join(script_dir, rel_path)

class CapturePacket ():
    def __init__(self, app_package_name, app_class_name, simulator_ip_port, capture_packet_mode):
        self._app_package_name = app_package_name
        self._app_class_name = app_class_name
        self._simulator_ip_port = simulator_ip_port
        self._capture_packet_mode = capture_packet_mode
        # 杀死本地adb进程。防止在循环多次执行过程中adb进程阻塞
        os.system('taskkill /f /im %s' % 'adb.exe')        
        # 建立连接
        time.sleep(1)
        command_str = "adb connect " + self._simulator_ip_port
        os.system(command_str)
        time.sleep(1)
        # 获取root权限 
        command_str = "adb -s " + self._simulator_ip_port + " root"
        os.system(command_str)
        time.sleep(1)
        # 自动获取安卓系统wlan0 IP地址
        command_str = "adb -s " + self._simulator_ip_port + " shell ip addr show wlan0"
        ip_str_src = re.sub(r"\s+", " ", os.popen(command_str).read().splitlines()[2])
        # print(ip_str_src)
        ip_str = ip_str_src.split(' ')[2].split('/')[0]
        self._ip = str(ip_str)
        print("System Wlan0 IP : ", self._ip)
        # 自动获取APP的userID
        command_str = "adb -s " + self._simulator_ip_port + " shell \"cat /data/system/packages.xml | grep " + self._app_package_name 
        userID = os.popen(command_str).read()
        for i in userID.split(' ') :
            if "userId" in i:
                userID = i.split('"')[1]
        self._userID = str(userID)
        print("APP User ID : ", self._userID)
        time.sleep(1)
        # 清空规则
        command_str = "adb -s " + self._simulator_ip_port + " shell iptables -F"
        os.popen(command_str)
        time.sleep(2)
        # 添加规则
        # 把OUTPUT数据包里面，用户ID是userID的流打上userID的流标签
        command_str = "adb -s " + self._simulator_ip_port + " shell iptables -A OUTPUT -m owner --uid-owner " + self._userID + " -j CONNMARK --set-mark "+ self._userID
        os.popen(command_str)
        time.sleep(2)
        # 把入站流中带有userID标签的数据包放到nflog消息池中，消息池的标号为userID
        command_str = "adb -s " + self._simulator_ip_port + " shell iptables -A INPUT -m connmark --mark " + self._userID + " -j NFLOG --nflog-group " + self._userID
        os.popen(command_str)
        time.sleep(2)
        # 把出站流中带有userID标签的数据包放到nflog消息池中，消息池的标号为userID
        command_str = "adb -s " + self._simulator_ip_port + " shell iptables -A OUTPUT -m connmark --mark "+ self._userID + " -j NFLOG --nflog-group " + self._userID
        os.popen(command_str)
        # 系统反应慢休息休息
        time.sleep(10)

    def __del__(self):
        # 删除Android内保存的pcap包
        command_str = "adb -s " + self._simulator_ip_port + " shell rm /sdcard/Pictures/*.pcap"
        os.popen(command_str)
        # 查看&&删除规则
        # command_str = "adb -s " + self._simulator_ip_port + " shell iptables -L"
        # return_str = os.popen(command_str).read()
        # print(return_str)
        command_str = "adb -s " + self._simulator_ip_port + " shell iptables -F"
        os.popen(command_str)
        # 关闭adb连接
        os.system('adb disconnect')
        # 杀死本地adb进程。防止在循环多次执行过程中adb进程阻塞
        os.system('taskkill /f /im %s' % 'adb.exe')
        
    def StartCapturePacket(self, pcap_file_ID, pcap_src_file_path, network_cord_name):   
        # 安进程抓包
        if self._capture_packet_mode & 0b0001:
            print("模拟器安进程抓包")
            command_str = "adb -s " + self._simulator_ip_port + " shell tcpdump -i nflog:" + self._userID +" -w  /sdcard/Pictures/simulator_uid_" + self._userID + "_" + str(pcap_file_ID) +".pcap &"
            os.popen(command_str)

        # 全系统抓包
        if self._capture_packet_mode & 0b0010:
            print("模拟器全系统抓包")  
            
            command_str = "adb -s " + self._simulator_ip_port + " shell tcpdump -w  /sdcard/Pictures/simulator_all_" + str(pcap_file_ID) +".pcap &"
            os.popen(command_str)

        if self._capture_packet_mode & 0b0100:
            print("Host机按模拟器IP抓包")  
            command = tshark_file_path + ' -i \"' + network_cord_name + '\"  -w ' + pcap_src_file_path + 'host_' + self._ip + '_' + str(pcap_file_ID)  + '.pcap -T text host ' + self._ip + ' > NUL'
            print(command)
            subprocess.Popen(command, shell=True)
        
        if self._capture_packet_mode & 0b1000:
            print("Host机全系统抓包")  
            command = tshark_file_path + ' -i \"' + network_cord_name + '\"  -w ' + pcap_src_file_path + 'host_all_' + str(pcap_file_ID)  + '.pcap' + ' -T text > NUL'
            print(command)
            subprocess.Popen(command, shell=True)

        return 
    
    def StopCapturePacket(self, pcap_file_ID, pcap_src_file_path):
        if self._capture_packet_mode & 0b0001:
            print("停止模拟器安进程抓包")
            command_str = "adb -s " + self._simulator_ip_port + " shell killall tcpdump"
            os.popen(command_str)
            # 查看保存的pcap包
            command_str = "adb -s " + self._simulator_ip_port + " shell ls -l /sdcard/Pictures/simulator_uid_" + self._userID + "_" + str(pcap_file_ID) +".pcap"
            return_str = os.popen(command_str).read()
            print(return_str)
            # pull pcap包 到本地
            command_str = "adb -s " + self._simulator_ip_port + " pull /sdcard/Pictures/simulator_uid_" + self._userID + "_" + str(pcap_file_ID) + ".pcap " + pcap_src_file_path
            return_str = os.popen(command_str).read()
            # 等待pcap包拷贝
            time.sleep(10)
            print(return_str)
            # 删除Android内保存的pcap包
            command_str = "adb -s " + self._simulator_ip_port + " shell rm /sdcard/Pictures/simulator_uid_"  + self._userID + "_" + str(pcap_file_ID) +".pcap"
            os.popen(command_str)          

        if self._capture_packet_mode & 0b0010:
            print("停止模拟器全系统抓包")  
            command_str = "adb -s " + self._simulator_ip_port + " shell killall tcpdump"
            os.popen(command_str)
            # 查看保存的pcap包
            command_str = "adb -s " + self._simulator_ip_port + " shell ls -l /sdcard/Pictures/simulator_all_" + str(pcap_file_ID) +".pcap &"
            return_str = os.popen(command_str).read()
            print(return_str)
            # pull pcap包 到本地
            command_str = "adb -s " + self._simulator_ip_port + " pull /sdcard/Pictures/simulator_all_" + str(pcap_file_ID) +".pcap " + pcap_src_file_path
            return_str = os.popen(command_str).read()
            # 等待pcap包拷贝
            time.sleep(10)
            print(return_str)
            # 删除Android内保存的pcap包
            command_str = "adb -s " + self._simulator_ip_port + " shell rm /sdcard/Pictures/simulator_all_" + str(pcap_file_ID) +".pcap"
            os.popen(command_str)  

        if self._capture_packet_mode & 0b0100:
            print("停止Host机按模拟器IP抓包")  
            process_name = 'tshark.exe'
            pid = os.popen('tasklist /fi "imagename eq %s"' % process_name).readlines()[3].split()[1]
            os.system('taskkill /F /PID %s' % pid)

        if self._capture_packet_mode & 0b1000:
            print("停止Host系统系统抓包")  
            process_name = 'tshark.exe'
            pid = os.popen('tasklist /fi "imagename eq %s"' % process_name).readlines()[3].split()[1]
            os.system('taskkill /F /PID %s' % pid)
        
        return

    def StartAPP(self):
        command_str = "adb -s " + self._simulator_ip_port + " shell am start -n " + self._app_package_name + self._app_class_name
        ret = os.popen(command_str).read()
        if not 'Starting' in ret:
            return -1
        return 1
    
    def StopAPP(self):
        command_str = "adb -s " + self._simulator_ip_port + " shell am force-stop " + self._app_package_name 
        print(os.popen(command_str).read())
        return
    
    def WgetFile(self, file_url):
        command_str = "adb -s " + self._simulator_ip_port + " shell wget " + file_url + " -O /sdcard/Pictures/wgetfile" 
        print(os.popen(command_str).read())
        return
    
    def CurlFile(self, file_url):
        command_str = "adb -s " + self._simulator_ip_port + " shell curl -k " + file_url
        os.popen(command_str).read()
        return 
    
    def GetPublicIP(self):
        command_str = "adb -s " + self._simulator_ip_port + " shell curl ifconfig.me/ip "
        string = os.popen(command_str).read()
        public_ip = string[string.rfind('\n') + 1:]
        return public_ip

    def GetLocation(self, ip):
        try:
            response = requests.get(f'https://ipinfo.io/{ip}/json')
            location_data = response.json()
            return location_data
        except Exception as e:
            print(f"Error: {e}")
            return None
        
    def SendDataByUdp(self, ip, port, file):
        # 拷贝文件到安卓模拟器
        for i in range(100):
            command_str = "adb -s " + self._simulator_ip_port + " shell \"echo 'Hello Word !!!' | busybox nc " + ip + " " + port + "\""
            string = os.popen(command_str).read()
        return string
    