python實現SSH 轉發

最近在做有關 python 的項目，工作的網絡環境如下：

    PC1是我所使用的開發機，處於公司內網中即都與公司10.239.58.1交換機相連，PC2有兩張網卡，第一張連接10.239.58.1，第二張連接192.168.122.1，操作系統為win8.1，PC3-PC6都是連入192.168.122.1，組成自己的局域網，而且操作系統均為Fedora LINUX。在PC1上開發的程序需要依賴PC3-PC6才能運行（依賴HTTP server 和 SSH登陸配置），所以整個流程是，我需要remote到PC2上，在PC2上進行開發測試，由於PC2是windows系統，不能同時多人remote，因此當我在使用PC2時其他同事若要想訪問PC3-PC5他們要去機房，那環境可想而知，或者我退出他們使用。這對於我來說太麻煩，多人共享一個設備很不方便，我決定做一個轉發工具，用於轉發HTTP 報文和對SSH的22端口進行轉發。由於剛剛接觸python，項目用的也是python，決定用python做這件事，還能再項目之余鍛煉一下。

    由於SSH使用的協議和HTTP使用的協議都是在TCP/IP協議之上，所以我就想到了在socket層上進行數據流轉發。代碼如下：

# -*- coding: utf-8 -*-
from socket import *
import threading
import Queue

# test in python 2.7

# a list of local port mapped to machine
ssh_server_list = [{"host": '10.239.58.102', "port": 8102, "ip": "192.168.122.102"},
                   {"host": '10.239.58.102', "port": 8103, "ip": "192.168.122.103"},
                   {"host": '10.239.58.102', "port": 8101, "ip": "192.168.122.101"},
                   {"host": '10.239.58.102', "port": 8104, "ip": "192.168.122.104"}]

# sync all the threads in each of tunnel
running_flag = []

# receive buffSize
buffer_size = 2048


def get_data_from_ssh_server(rev_msg, tcp_socket, flag):
    """
    :param rev_msg: a queue buffer of message need to be send to SSH client
    :param tcp_socket: instance of socket used for sending data
    :param flag: control this function
    :return: null
    """
    while running_flag[flag]:
        data = tcp_socket.recv(buffer_size)
        if len(data):
            rev_msg.put(str(data))
        else:
            running_flag[flag] = False


def send_data_to_ssh_client(rev_msg, tcp_socket, flag):
    """
    :param rev_msg: a queue buffer of message need to be send to SSH client
    :param tcp_socket: instance of socket used for sending data
    :param flag: control this function
    :return: null
    """
    while running_flag[flag]:
        try:
            data = rev_msg.get(timeout=10)
            data = tcp_socket.send(str(data))
        except:
            pass


def get_data_from_ssh_client(send_msg, tcp_socket, flag):
    """
    :param send_msg: a queue buffer of message need to be send to SSH server in each machine
    :param tcp_socket: instance of socket used for sending data
    :param flag: control this function
    :return: null
    """
    while running_flag[flag]:
        data = tcp_socket.recv(buffer_size)
        if len(data):
            send_msg.put(str(data))
        else:
            running_flag[flag] = False


def send_data_to_ssh_server(send_msg, tcp_socket, flag):
    """
    :param send_msg: a queue buffer of message need to be send to SSH server in each machine
    :param tcp_socket: instance of socket used for sending data
    :param flag: control this function
    :return: null
    """
    while running_flag[flag]:
        try:
            data = send_msg.get(timeout=10)
            data = tcp_socket.send(str(data))
        except:
            pass


def handle_connections(host, ip, port):
    """
    :param host: local ip
    :param ip: which machine the data will be forwarded
    :param port: local port
    :return: null
    """
    ssh_client_socket = socket(AF_INET, SOCK_STREAM)
    ssh_client_socket.bind((host, port))

    # listen 10 client
    ssh_client_socket.listen(10)
    while True:
        ssh_client_side, address = ssh_client_socket.accept()

        # two queue for keeping data from SSH client and SSH server
        buffer_send = Queue.Queue()
        buffer_rev = Queue.Queue()

        ssh_server_side = socket(AF_INET, SOCK_STREAM)
        ssh_server_side.connect((ip, 22))

        flag = True

        running_flag.append(flag)

        rev1 = threading.Thread(target=get_data_from_ssh_server,
                                args=(buffer_rev, ssh_server_side, len(running_flag) - 1))
        rev2 = threading.Thread(target=send_data_to_ssh_client,
                                args=(buffer_rev, ssh_client_side, len(running_flag) - 1))

        send1 = threading.Thread(target=get_data_from_ssh_client,
                                 args=(buffer_send, ssh_client_side, len(running_flag) - 1))
        send2 = threading.Thread(target=send_data_to_ssh_server,
                                 args=(buffer_send, ssh_server_side, len(running_flag) - 1))

        rev1.start()
        rev2.start()
        send1.start()
        send2.start()


if __name__ == "__main__":
    print("start SSH forward server")

    thread_pool = []
    for i in ssh_server_list:
        print("ssh mapping " + i["host"] + ":" + str(i["port"]) + " => " + i["ip"] + ":22")
        t = threading.Thread(target=handle_connections, args=(i["host"], i["ip"], i["port"]))
        thread_pool.append(t)
        t.start()

    print("initialize SSH forward server done")

寫完按這個小工具 頓時感覺Python還是很強大的！