codewars/books/python_black/scanner.py

import ipaddress
import os
import socket
import struct
import sys
import threading
import time


# сканируем подсеть
SUBNET = '192.168.31.0/24'
# волшебная строка, которую мы будем искать в ICMP-ответах
MESSAGE = 'PYTHONRULES!'

class IP:
    def __init__(self, buff):
        header = struct.unpack('<BBHHHBBH4s4s', buff)
        self.ver = header[0] >> 4
        self.ihl = header[0] & 0xF

        self.tos = header[1]
        self.len = header[2]
        self.id = header[3]
        self.offset = header[4]
        self.ttl = header[5]
        self.protocol_num = header[6]
        self.sum = header[7]
        self.src = header[8]
        self.dst = header[9]

        # IP-адруса, понятные человек
        self.src_address = ipaddress.ip_address(self.src)
        self.dst_address = ipaddress.ip_address(self.dst)

        # сопоставляем константы протоколов с их назначением
        self.protocol_map = {1: "ICMP", 6: "TCP", 17: "UDP"}
        try:
            self.protocol = self.protocol_map[self.protocol_num]
        except Exception as e:
            print('%s No protocol for %s' % (e, self.protocol_num))
            self.protocol = str(self.protocol_num)


class ICMP:
    def __init__(self, buff):
        header = struct.unpack('<BBHHH', buff)
        self.type = header[0]
        self.code = header[1]
        self.sum = header[2]
        self.id = header[3]
        self.seq = header[4]


# Добавляет в UDP-датаграммы наше волшебное сообщение
def udp_sender():
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sender:
        for i in range(10):
            # sender.sendto(bytes(MESSAGE, 'utf8'), ('192.168.31.100', 65212))
            sender.sendto(bytes(MESSAGE, 'utf8'), ('8.8.8.8', 65212))
            time.sleep(0.3)
        # for ip in ipaddress.ip_network(SUBNET).hosts():
            # print(bytes(MESSAGE, 'utf8'), (str(ip), 65212))
            # sender.sendto(bytes(MESSAGE, 'utf8'), (str(ip), 65212))


class Scanner:
    def __init__(self, host):
        self.host = host
        if os.name == 'nt':
            socket_protocol = socket.IPPROTO_IP
        else:
            socket_protocol = socket.IPPROTO_ICMP
        
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket_protocol)
        self.socket.bind((host, 0))
        self.socket.setsockopt(socket.IPPROTO_IP, socket.IP_HDRINCL, 1)

        if os.name == 'nt':
            self.socket.ioctl(socket.SIO_RCVALL, socket.RCVALL_ON)



    def sniff(self):
        hosts_up = set([f'{str(self.host)} *'])
        try:
            while True:
                # читаем пакет
                raw_buffer = self.socket.recvfrom(65535)[0]
                # создаем IP-заголовок из первых 20 байт
                ip_header = IP(raw_buffer[0:20])
                # нас интересует ICMP
                print(raw_buffer)
                print(ip_header.src_address, ip_header.dst_address)
                if ip_header.protocol == 'ICMP':
                    print('ICMP')
                    offset = ip_header.ihl * 4
                    buf = raw_buffer[offset:offset + 8]
                    # print(raw_buffer)
                    icmp_header = ICMP(buf)
                    # ищем тип и код 3
                    print(icmp_header.code, icmp_header.type)
                    if icmp_header.code == 3 and icmp_header.type == 3:
                        print('!')
                        if ipaddress.ip_address(ip_header.src_address) in ipaddress.IPv4Network(SUBNET):
                            # проверяем содержит ли буфер наше волшебное сообщение
                            print(raw_buffer[len(raw_buffer) - len(MESSAGE):])
                            if raw_buffer[len(raw_buffer) - len(MESSAGE):] == bytes(MESSAGE, 'utf8'):
                                tgt = str(ip_header.src_address)
                                print(ip_header.src_address, ip_header.dst_address)
                                if tgt != self.host and tgt not in hosts_up:
                                    hosts_up.add(str(ip_header.dst_address))
                                    print(f'Host Up: {tgt}')
        # обрабатываем Ctrl+C
        except KeyboardInterrupt:
            # если мы в Windows, выключаем неизбирательный режим
            if os.name == 'nt':
                self.socket.ioctl(socket.SIO_RCVALL, socket.RCVALL_OFF)
            
            print('\nUser interrupted')
            
            if hosts_up:
                print(f'\n\nSummary: Hosts up on {SUBNET}')
                for host in sorted(hosts_up):
                    print(f'{host}')
                print('')
                sys.exit()



if __name__ == '__main__':
    if len(sys.argv) == 2:
        host = sys.argv[1]
        s = Scanner(host)
        # time.sleep(1)
        t = threading.Thread(target=udp_sender)
        t.start()
        s.sniff()