zrok/sdk/python/examples/pastebin/pastebin.py
2024-03-11 15:36:24 -04:00

104 lines
2.9 KiB
Python
Executable File

#!python3
import argparse
import atexit
import sys
import os
import zrok
import zrok.listener
import zrok.dialer
from zrok.model import AccessRequest, ShareRequest
import signal
import threading
exit_signal = threading.Event()
def signal_handler(signum, frame):
print("\nCtrl-C detected. Next connection will close server")
exit_signal.set()
class copyto:
def handle(self, *args, **kwargs):
root = zrok.environment.root.Load()
try:
shr = zrok.share.CreateShare(root=root, request=ShareRequest(
BackendMode=zrok.model.TCP_TUNNEL_BACKEND_MODE,
ShareMode=zrok.model.PRIVATE_SHARE_MODE,
Target="pastebin"
))
except Exception as e:
print("unable to create share", e)
sys.exit(1)
def removeShare():
try:
zrok.share.DeleteShare(root, shr)
except Exception as e:
print("unable to delete share", e)
sys.exit(1)
atexit.register(removeShare)
data = self.loadData()
print("access your pastebin using 'pastebin.py pastefrom " + shr.Token + "'")
with zrok.listener.Listener(shr.Token, root) as server:
while not exit_signal.is_set():
conn, peer = server.accept()
with conn:
conn.sendall(data.encode('utf-8'))
print("Server stopped.")
def loadData(self):
if not os.isatty(sys.stdin.fileno()):
return sys.stdin.read()
else:
raise Exception("'copyto' requires input from stdin; direct your paste buffer into stdin")
def pastefrom(options):
root = zrok.environment.root.Load()
try:
acc = zrok.access.CreateAccess(root=root, request=AccessRequest(
ShareToken=options.shrToken,
))
except Exception as e:
print("unable to create access", e)
sys.exit(1)
def removeAccess():
try:
zrok.access.DeleteAccess(root, acc)
except Exception as e:
print("unable to delete access", e)
sys.exit(1)
atexit.register(removeAccess)
client = zrok.dialer.Dialer(options.shrToken, root)
data = client.recv(1024)
print(data.decode('utf-8'))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers()
subparsers.required = True
c = copyto()
parser_copyto = subparsers.add_parser('copyto')
parser_copyto.set_defaults(func=c.handle)
parser_pastefrom = subparsers.add_parser('pastefrom')
parser_pastefrom.set_defaults(func=pastefrom)
parser_pastefrom.add_argument("shrToken")
options = parser.parse_args()
signal.signal(signal.SIGINT, signal_handler)
# Create a separate thread to run the server so we can respond to ctrl-c when in 'accept'
server_thread = threading.Thread(target=options.func, args=[options])
server_thread.start()
server_thread.join()