CLOSE_WAIT connection killer for ROS

kill_close_wait.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright: Yuki Furuta <furushchev@jsk.imi.i.u-tokyo.ac.jp>

import socket
import struct
import traceback
try:
    import psutil
except:
    print "Please execute 'sudo apt-get install python-psutil'"
    exit(1)

def tcp_checksum(msg):
    s = 0

    # loop taking 2 characters at a time
    for i in range(0, len(msg), 2):
        w = ord(msg[i]) + (ord(msg[i+1]) << 8 )
        s = s + w

    s = (s>>16) + (s & 0xffff);
    s = s + (s >> 16);

    #complement and mask to 4 byte short
    s = ~s & 0xffff
    return s

packet_id = 1
def create_tcp_packet_raw(src_addr, src_port, dst_addr, dst_port):
    global packet_id
    ip_ihl = 5
    ip_ver = 4
    ip_tos = 0
    ip_tot_len = 0
    ip_id = packet_id
    packet_id += 1
    ip_frag_off = 0
    ip_ttl = 64
    ip_proto = socket.IPPROTO_TCP
    ip_check = 0
    ip_saddr = socket.inet_aton(src_addr)
    ip_daddr = socket.inet_aton(dst_addr)
    ip_ihl_ver = (ip_ver << 4) + ip_ihl
    ip_header = struct.pack('!BBHHHBBH4s4s',
                            ip_ihl_ver, ip_tos, ip_tot_len, ip_id, ip_frag_off, ip_ttl,
                            ip_proto, ip_check, ip_saddr, ip_daddr)

    tcp_src = src_port
    tcp_dst = dst_port
    tcp_seq = 10
    tcp_ack_seq = 0
    tcp_doff = 5
    tcp_fin = 0
    tcp_syn = 0
    tcp_rst = 0
    tcp_psh = 0
    tcp_ack = 1
    tcp_urg = 0
    tcp_window = socket.htons(65535)
    tcp_check = 0
    tcp_urg_ptr = 0
    tcp_offset_res = (tcp_doff << 4) + 0
    tcp_flags = tcp_fin + (tcp_syn << 1) + (tcp_rst << 2) + (tcp_psh << 3) + (tcp_ack << 4) + (tcp_urg << 5)
    tcp_header = struct.pack('!HHLLBBHHH', tcp_src, tcp_dst, tcp_seq, tcp_ack_seq, tcp_offset_res, tcp_flags, tcp_window, tcp_check, tcp_urg_ptr)

    tcp_src_addr = socket.inet_aton(src_addr)
    tcp_dst_addr = socket.inet_aton(dst_addr)
    placeholder = 0
    protocol = socket.IPPROTO_TCP
    tcp_length = len(tcp_header)

    psh = struct.pack('!4s4sBBH', tcp_src_addr, tcp_dst_addr, placeholder, protocol, tcp_length)
    psh = psh + tcp_header

    tcp_check = tcp_checksum(psh)

    tcp_header = struct.pack('!HHLLBBH', tcp_src, tcp_dst, tcp_seq, tcp_ack_seq, tcp_offset_res, tcp_flags, tcp_window)
    tcp_header += struct.pack('H', tcp_check) + struct.pack('!H', tcp_urg_ptr)

    packet = ip_header + tcp_header
    return packet

def main():
    # find rosmaster
    rosmaster = None
    for p in psutil.process_iter():
        if p.name == "rosmaster":
            rosmaster = p
    if not rosmaster:
        raise "rosmaster not found"

    # find CLOSE_WAIT connections
    close_waits = filter(lambda c: c.family == socket.AF_INET and c.status == "CLOSE_WAIT",
                         rosmaster.get_connections("tcp"))
    print "Found %d CLOSE_WAIT Connections" % len(close_waits)

    # create raw packet socket
    s = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_RAW)
    s.setsockopt(socket.IPPROTO_IP, socket.IP_HDRINCL, 1)

    # send ack packet to close waiting connections
    for c in close_waits:
        dst_addr, dst_port = c.local_address
        src_addr, src_port = c.remote_address
        try:
            packet = create_tcp_packet_raw(src_addr, src_port, dst_addr, dst_port)
            s.sendto(packet, (dst_addr, dst_port))
        except Exception as e:
            print str(e), traceback.format_exc()
    s.close()

if __name__ == '__main__':
    import os
    if os.getuid() != 0:
        print "Please run with sudo!"
        exit(1)
    else:
        main()