core.py

#!/usr/bin/env python
# coding: utf8

import socket
import os
import sys
import random
import struct
import time
from collections import Counter, defaultdict

class Tracer:
    def __init__(self, family=socket.AF_INET):
        self.family = family
        self.ipv6 = (family == socket.AF_INET6)
        self.batch_size = 100
        self.max_retry = 10
        self.timeout = 1
        
        if self.ipv6:
            self.proto, self.ttl_opt = socket.IPPROTO_ICMPV6, socket.IPV6_UNICAST_HOPS
            self.type_echo, self.type_timxed, self.type_reply = 128, 3, 129
        else:
            self.proto, self.ttl_opt = socket.IPPROTO_ICMP, socket.IP_TTL
            self.type_echo, self.type_timxed, self.type_reply = 8, 11, 0

        self.sock = socket.socket(self.family, socket.SOCK_RAW, self.proto)
        self.sock.settimeout(self.timeout)
        
        self.max_ttl = defaultdict(lambda: 32)
        self.echo_map = {} # {id: (target, ttl, ts)}
        self.rtt_map = defaultdict(list)
        self.result = defaultdict(dict)
        self.sent_counts = Counter()
        self.recv_counts = Counter()
        
        self.in_flight = set()
        self.retries = Counter()
        self.running = True
        self.on_pong = lambda t, target, hop, ttl, rtt, is_target: None

    def _get_checksum(self, msg):
        s = 0
        for i in range(0, len(msg), 2):
            w = msg[i] + (msg[i+1] << 8)
            s = (s + w); s = (s & 0xffff) + (s >> 16)
        return ~s & 0xffff

    def create_packet(self, icmp_id, ts):
        header = struct.pack('!BBHHH', self.type_echo, 0, 0, icmp_id, 1)
        payload = struct.pack('!Q', ts) + b'trpy'
        if self.ipv6: return header + payload
        msg = bytearray(header + payload)
        cs = self._get_checksum(msg)
        struct.pack_into('<H', msg, 2, cs)
        return msg

    def ping(self, ip, ttl, max_counts=None):
        if max_counts and self.sent_counts[(ip, ttl)] >= max_counts: return
        self.sock.setsockopt(socket.IPPROTO_IP if not self.ipv6 else socket.IPPROTO_IPV6, self.ttl_opt, ttl)
        icmp_id = random.randint(1000, 65000)
        ts = int(time.time() * 1000000)
        self.echo_map[icmp_id] = (ip, ttl, ts)
        self.sent_counts[(ip, ttl)] += 1
        try: self.sock.sendto(self.create_packet(icmp_id, ts), (ip, 0))
        except: pass

    def tick(self, queue, counts=None):
        sent = 0
        for key in list(self.in_flight):
            if self.retries[key] > 0:
                if not counts or self.sent_counts[key] < counts:
                    self.ping(*key, max_counts=counts)
                    sent += 1
                self.retries[key] -= 1
            else:
                self.in_flight.discard(key); self.retries.pop(key, None)
                target, ttl = key
                if ttl not in self.result[target]: self.result[target][ttl] = []
            if sent >= self.batch_size: return

        while sent < self.batch_size:
            try:
                target, ttl = next(queue)
                if ttl > self.max_ttl[target]: continue
                self.ping(target, ttl, max_counts=counts)
                self.retries[(target, ttl)] = self.max_retry
                self.in_flight.add((target, ttl)); sent += 1
            except StopIteration:
                self.running = False; return

    def run(self, target_ips, counts=None):
        def q_gen():
            c = 0
            while True:
                for ip in target_ips:
                    for ttl in range(1, 40): yield (ip, ttl)
                c += 1
                if counts and c >= counts: break
        queue = q_gen()
        self.running = True
        self.tick(queue, counts=counts)
        while self.running or self.in_flight:
            try:
                data, addr = self.sock.recvfrom(1280)
                self.on_data(data, addr[0], counts=counts)
            except socket.timeout: self.tick(queue, counts=counts)
            except KeyboardInterrupt: break
        return self.result

    def on_data(self, data, hop_ip, counts=None):
        if self.ipv6:
            if len(data) < 6: return
            icmp_type, icmp_code, _, icmp_id = struct.unpack('!BBHH', data[:6])
            body = data
        else:
            ip_hl = (data[0] & 0x0F) * 4
            if len(data) < ip_hl + 6: return
            body = data[ip_hl:]; icmp_type, icmp_code, _, icmp_id = struct.unpack('!BBHH', body[:6])

        if icmp_type == self.type_timxed:
            if len(body) < 8: return
            orig_ip_start = 8 + (40 if self.ipv6 else (body[8] & 0x0F) * 4)
            if len(body) < orig_ip_start + 6: return
            _, _, _, icmp_id = struct.unpack('!BBHH', body[orig_ip_start : orig_ip_start + 6])
        elif icmp_type != self.type_reply: return

        if icmp_id in self.echo_map:
            target_ip, ttl, send_ts = self.echo_map.pop(icmp_id)
            if ttl > self.max_ttl[target_ip]: return
            rtt = (int(time.time() * 1000000) - send_ts) / 1000.0
            
            is_target = (hop_ip == target_ip or icmp_type == self.type_reply)
            if is_target:
                if ttl < self.max_ttl[target_ip]:
                    self.max_ttl[target_ip] = ttl
                    for t in range(ttl + 1, 40):
                        self.in_flight.discard((target_ip, t)); self.retries.pop((target_ip, t), None)
            
            if ttl not in self.result[target_ip]: self.result[target_ip][ttl] = []
            if hop_ip not in self.result[target_ip][ttl]: self.result[target_ip][ttl].append(hop_ip)
            self.rtt_map[(target_ip, ttl, hop_ip)].append(rtt)
            self.recv_counts[(target_ip, ttl)] += 1
            if counts and self.recv_counts[(target_ip, ttl)] >= counts:
                self.in_flight.discard((target_ip, ttl)); self.retries.pop((target_ip, ttl), None)
            self.on_pong(self, target_ip, hop_ip, ttl, rtt, is_target)