from base import BaseServer

from gevent.server import StreamServer

from gevent import monkey
monkey.patch_all()

import socket
import sys
import threading
import time
import os

global stored_files, serv_arch, serv_compile
stored_files = {}

class RetrieveServer(BaseServer):

	def __init__(self, parent):
		BaseServer.__init__(self, parent)
		self.directory = "../bins"

	def send_file(self, arch, sock):
		file_name = str(arch)

		dirlist = os.listdir(self.directory) # anti lfi

		if not file_name in dirlist:
			return

		f = open("%s/%s" % (self.directory, file_name), "rb")
		data = f.read()
		f.close()
		f2 = open("%s/%s" % (self.directory, file_name), "rb")
		i = len(data)
		while True:
			g = f2.read(1000)
			n = sock.send(bytearray(g))
			if n != len(g):
				sock.send(bytearray(g)[n:])
			if n > 0:
				i -= n
			if i == 0 or n == 0:
				break
		f2.close()

	def start(self, port, args):
		
		listen = ""
		argc = len(args)
		self.directory = "../bins"

		if argc >= 1:
			self.directory = args[0]
		if argc >= 2:
			listen = args[1]
		
		def handle(sock, addr):
			try:
				arch = sock.recv(20).strip()
				self.send_file(arch, sock)
				print("[RETRIEVE-SERVER] Sent " + arch + " to " + addr[0])
				sock.close()
			except Exception as e:
				print("[RETRIEVE-SERVER] Client error - " + str(addr) + " - " + str(e))
				sock.close()
		server = StreamServer((listen, port), handle)
		server.start()
		return server

	def default_port(self):
		return 25566

	def default_args(self):
		return "../bins 0.0.0.0"

	def usage(self):
		return "<directory> <listen ip>"

def main(argv):
	argv = argv[1:]
	argc = len(argv)
	dir = "../bins"
	port = 25566

	if argc == 1:
		port = argv[0]
	elif argc == 2:
		port = argv[0]
		dir = argv[1]
	
	server = RetrieveServer(None)
	server.directory = dir
	server.start(port, [])

if __name__ == '__main__':
	main(sys.argv)