client.py

import os
import socket
import sys
import threading
import time

import auth
from config import (HOST, PORT, TUN_NAME, TUN_CLIENT_ADDR, MAX_PACKET,
                    FLAG_HEARTBEAT)
from protocol import send_msg, recv_msg
from tun import open_tun, configure_tun, ip_src, ip_dst

HEARTBEAT_INTERVAL = 15
running = True


def tun_to_socket(tun_fd, sock):
    global running
    while running:
        try:
            packet = os.read(tun_fd, MAX_PACKET)
        except OSError:
            break
        if not packet:
            break
        try:
            send_msg(sock, packet)
        except OSError:
            break
    running = False


def heartbeat(sock):
    while running:
        time.sleep(HEARTBEAT_INTERVAL)
        try:
            send_msg(sock, b"", 0, FLAG_HEARTBEAT)
        except OSError:
            break


def socket_to_tun(tun_fd, sock):
    global running
    while running:
        try:
            msg = recv_msg(sock)
        except OSError:
            break
        if msg is None:
            print("server closed the connection")
            break

        payload, session_id, flags = msg
        if flags == FLAG_HEARTBEAT:
            continue
        if len(payload) < 20:
            continue

        os.write(tun_fd, payload)
    running = False


def main():
    if len(sys.argv) < 3:
        sys.exit("usage: sudo python3 client.py <username> <password>")

    username, password = sys.argv[1], sys.argv[2]

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((HOST, PORT))
        print(f"connected to {HOST}:{PORT}")

        ok, message = auth.client_authenticate(s, username, password)
        print(message)
        if not ok:
            return

        tun_fd = open_tun(TUN_NAME)
        configure_tun(TUN_NAME, TUN_CLIENT_ADDR)
        print(f"{TUN_NAME} up with {TUN_CLIENT_ADDR}")
        print(f"now run: sudo ./setup_client.sh {HOST}")

        threading.Thread(target=tun_to_socket, args=(tun_fd, s), daemon=True).start()
        threading.Thread(target=heartbeat, args=(s,), daemon=True).start()
        socket_to_tun(tun_fd, s)


if __name__ == "__main__":
    if os.geteuid() != 0:
        sys.exit("run with sudo")
    try:
        main()
    except ConnectionRefusedError:
        print("connection refused - start server.py first")
    except KeyboardInterrupt:
        print("stopped by user")