whitelisted_socket_guard.py

# basic whitelisting for socket.socket.connect.

import socket

CONNECTION_WHITELIST = [('localhost', 80)]

unguarded_connect = socket.socket.connect

def guarded_connect(self, *args, **kwargs):
    con = args[0]
    if con not in CONNECTION_WHITELIST:
        raise Exception(f"Not allowed! {con} not in whitelist")

    return unguarded_connect(self, *args, **kwargs)

socket.socket.connect = guarded_connect


# this works
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    s.connect(("localhost", 80))
    s.sendall(b'hello world')
    data = s.recv(1024)

# this doesn't
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    s.connect(("google.com", 80))
    s.sendall(b'hello world')
    data = s.recv(1024)