fstream/client/fstream.py

#!/usr/bin/env python3
import sys
import os
p = os.path.join
pUp = os.path.dirname
s = False
if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
	s = os.path.realpath(sys.executable)
else:
	s = os.path.realpath(__file__)
sp = pUp(s)

# script start
import subprocess
import socket
import time

def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)
def defGet(tbl,key,defv):
	if not key in tbl: return defv
	return tbl[key]

bufferSize = int(defGet(os.environ,"fstream_net_buffer","4096")) # max buffer size in bytes for receiving data, lower values shouldn't reduce the delay
bufferSizeStdin = int(defGet(os.environ,"fstream_stdin_buffer","128")) # min buffer size for buffer, lower values DO reduce delay but raise CPU usage
timeout = float(defGet(os.environ,"fstream_net_timeout","15")) # timeout in seconds
useSSL = False
sslIgnoreCert = False
if defGet(os.environ,"fstream_ssl","0") == "1": useSSL = True
if defGet(os.environ,"fstream_ssl_ignoreCert","0") == "1": sslIgnoreCert = True

connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
unbufferedStdout = os.fdopen(sys.stdout.fileno(),"wb",0) # Make unbuffered stdout

def listToCommand(lst):
	cmd = ""
	for arg in lst:
		arg = arg.replace("\\","\\\\")
		arg = arg.replace(",","\\,")
		cmd += arg + ","
	
	return cmd[:-1]

def recv(conn,l):
	start = time.process_time()
	timeo = conn.gettimeout()
	bytes = b""
	while l > 0:
		b = conn.recv(l)
		if b == b"": raise ConnectionResetError
		if time.process_time() - start > timeo: raise TimeoutError
		bytes += b
		l -= len(b)
	return bytes

def getResponse(connection,maxLength = 0):
	data = b''
	data = recv(connection,4)
	if not data: return False
	nul = recv(connection,1)
	if not nul: return False
	if nul != b"\x00": return False
	requestLength = int.from_bytes(data,"big")
	if maxLength != 0 and requestLength > maxLength: return False
	return recv(connection,requestLength)

def sendResponse(connection,data):
	connection.sendall(len(data).to_bytes(4,"big") + b"\x00" + data)

def stringToAddressTuple(addr):
	rtn = addr.rsplit(":",1)
	rtn[1] = int(rtn[1])
	rtn = tuple(rtn)
	return rtn

def main():
	global serverAddr, useSSL, sslIgnoreCert, connection
	serverAddr = stringToAddressTuple(sys.argv[1])
	if useSSL: import ssl
	
	serverEnv = [sys.argv[2]]
	for var in os.environ:
		if not var.startswith("fstream_arg_"): continue
		serverEnv.append(var.replace("fstream_arg_","",1)+ "=" +os.environ[var])
	serverEnv = listToCommand(serverEnv)
	
	eprint("Connecting to server...")
	connection.settimeout(timeout)
	connection.connect(serverAddr)
	if useSSL:
		eprint("Performing SSL handshake...")
		if sys.version_info >= (3,10):
			proto = ssl.PROTOCOL_TLS_CLIENT
		else:
			proto = ssl.PROTOCOL_TLS
		
		ctx = ssl.SSLContext(proto)
		if sslIgnoreCert:
			ctx.check_hostname = False
			ctx.verify_mode = ssl.CERT_NONE
		
		connection = ctx.wrap_socket(
			connection,
			server_side = False
		)
	
	eprint("Sending payload...")
	connection.sendall("\n\n".encode("ascii"))
	sendResponse(connection,serverEnv.encode("utf-8"))
	
	if sys.argv[2] == "token":
		try:
			token = b""
			eprint("Receiving token...")
			while True:
				data = connection.recv(bufferSize)
				if data != b"":
					token += data
					continue
				
				if token == b"":
					eprint("Connection closed: failure")
					os.exit(1)
				else:
					eprint("Connection closed: success")
					unbufferedStdout.write(token)
				return
		except:
			connection.close()
			raise
	
	if sys.argv[2] == "watch":
		try:
			eprint("Receiving data...")
			while True:
				data = connection.recv(bufferSize)
				if data == b"":
					eprint("Connection closed.")
					return
				unbufferedStdout.write(data)
		except:
			connection.close()
			raise
	
	if sys.argv[2] == "broadcast":
		try:
			eprint("Sending data...")
			while True:
				data = sys.stdin.buffer.read(bufferSizeStdin)
				if data == b"":
					return
				connection.sendall(data)
		except:
			connection.close()
			raise

if __name__ == '__main__':
	main()