python学习之socket应用

环境说明

本文的代码运行环境：
windows 10，python 3.7（代码兼容python 2）

实现的目标：
1、实现nslookup 命令，查询域名对应的ip地址
2、实现ping命令，查询目标ip是否在线

代码实现

nslookup命令实现

import socket
import sys


def nslookup(domain):
    result = socket.getaddrinfo(domain,53,socket.AF_INET,socket.SOCK_DGRAM,17,0)
    #print(result)
    num=len(result)
    #print(num)
    for a in range(num):
        i = result[a][4][0]
        print(i)


if __name__ == '__main__':
    try:
        domain = sys.argv[1]
    except:
        print('Usage: %s www.baidu.com' % sys.argv[0])
        exit()
    nslookup(domain)

ping命令实现（icmp协议数据包模拟）

# coding=utf-8
import os
import sys
import socket
import struct
import select

try:
    from time import perf_counter
except:
    from time import clock as perf_counter

ICMP_ECHO_REQUEST = 8

def receive_ping(my_socket, ID, timeout):
    """
    receive the ping from the socket
    """
    start_time = timeout
    while True:
        start_select = perf_counter()
        # select.select(rlist, wlist, xlist[, timeout])
        # wait until ready for read / write / exceptional condition
        # The return value is a triple of lists
        what_ready = select.select([my_socket], [], [], start_time)
        how_long = perf_counter() - start_select
        if what_ready[0] == []: #timeout
            return None
        time_received = perf_counter()
        # socket.recvfrom(bufsize[, flags])
        # The return value is a pair (string, address)
        rec_packet, addr = my_socket.recvfrom(1024)
        icmp_header = rec_packet[20:28]
        ip_type, code, checksum, packet_ID, sequence = struct.unpack("bbHHh", icmp_header)
        if ip_type != 8 and packet_ID == ID: # ip_type should be 0
            byte_in_double = struct.calcsize("d")
            time_sent = struct.unpack("d", rec_packet[28:28 + byte_in_double])[0]
            return time_received - time_sent
        start_time = start_time - how_long
        if start_time <= 0:
            return None


def get_checksum(source):
    """
    return the checksum of source
    the sum of 16-bit binary one's complement
    """
    #print(type(source)) # in py2, str; in py3, bytes
    checksum = 0
    count = int((len(source) // 2) * 2)
    i = 0
    while i < count:
        if sys.version.startswith('3.'):
            temp = source[i + 1] * 256 + source[i] # 256 = 2^8
        else:
            temp = ord(source[i + 1]) * 256 + ord(source[i]) # 256 = 2^8
        checksum += temp
        checksum = checksum & 0xffffffff # 4,294,967,296 (2^32)
        i = i + 2
    if i < len(source):
        if sys.version.startswith('3.'):
            checksum += source[len(source) - 1]
        else:
            checksum += ord(source[len(source) - 1])
        checksum = checksum & 0xffffffff
    # 32-bit to 16-bit
    checksum = (checksum >> 16) + (checksum & 0xffff)
    checksum = checksum + (checksum >> 16)
    answer = ~checksum
    answer = answer & 0xffff
    # why? ans[9:16 1:8]
    answer = answer >> 8 | (answer << 8 & 0xff00)
    return answer



def send_ping(my_socket, ip_addr, ID):
    """
    send ping to the given ip address
    """
    ip = socket.gethostbyname(ip_addr)
    # Header is type (8), code (8), checksum (16), id (16), sequence (16)
    my_checksum = 0
    # Make a dummy heder with a 0 checksum
    # b:signed char, h:short 2, H:unsigned short 2
    header = struct.pack('bbHHh', ICMP_ECHO_REQUEST, 0, my_checksum, ID, 1)
    # struct.calcsize(fmt)
    # Return the size of the struct corresponding to the given format.
    byte_in_double = struct.calcsize("d") # C type: double, 8
    data = (192 - byte_in_double) * b"P" # any char is OK, any length is OK
    tt = perf_counter()
    data = struct.pack("d", tt) + data
    my_checksum = get_checksum(header + data)
    header = struct.pack("bbHHh", ICMP_ECHO_REQUEST, 0, socket.htons(my_checksum), ID, 1)
    packet = header + data
    my_socket.sendto(packet, (ip, 80)) # it seems that 0~65535 is OK (port?)


def ping_once(ip_addr, timeout):
    """
    return either delay (in second) or none on timeout.
    """
    icmp = socket.getprotobyname('icmp')
    try:
        my_socket = socket.socket(socket.AF_INET, socket.SOCK_RAW, icmp)
    except socket.error:
        raise
    # Return the current process id.
    # int: 0xFFFF = -1, unsigned int: 65535
    my_ID = os.getpid() & 0xFFFF
    send_ping(my_socket, ip_addr, my_ID)
    delay = receive_ping(my_socket, my_ID, timeout)
    my_socket.close()
    return delay


def icmp_ping(ip_addr, timeout = 2, count = 1):
    """
    send ping to ip_addr for count times with the given timeout
    """
    for i in range(count):
        print('ping ' + ip_addr)
        try:
            delay = ping_once(ip_addr, timeout)
        except socket.gaierror as e:
            print("failed. (socket error: '%s')" % e[1])
            break
        if delay == None:
            print('failed. (timeout within %s second.)' % timeout)
        else:
            print('get reply in %0.4f ms' % (delay * 1000))



if __name__ == '__main__':
    try:
        ip_addr = sys.argv[1]
    except:
        print('Usage: %s 1.2.3.4' % sys.argv[0])
        exit()
    icmp_ping(ip_addr)

参考文献：
ICMP协议Ping方法的Python实现解析