#!/usr/bin/python

import socket, struct, hashlib, sys, getopt, time

"""Simple Socket File Transfer,
originally by Whitsoft Development (http://www.whitsoftdev.com/ssft/)

Protocol:
	Greeting: active sends greeting to passive, passive responds with protocol id.
	Greeting packet:
		1) dword: id, "SSFT"
		2) dword: version, 0x10000
		3) byte: direction, 0 = send, 1 = recv
	Transfer ("=>": to receiver, "<=": to sender):
		1) => word: file name length
		2) => wide-chars: file name
		3) => long long: file size
		4) <= long long: file size at receiver
		5) if 4) != 0:
			5a) => file hash of the first 4) bytes
			5b) <= byte: continue transfer at 4) (1 / 0)
		6) => data (in 1024b chunks, hash calculated twice for each o_O)
		7) => hash

contact: ksooda at gmail com.
Finnish homepage: http://sooda.dy.fi

2008/02/23: slighty adjusted by mattijs to support ipv6
contact: mattijs at nxdomain dot nl
"""


class SsftSock:
	"""Base class for sender and receiver"""
	# the encoding used in command line parameters
	ENCODING = "iso-8859-1"
	PROTOCOL_ID = ord("S") | (ord("S") << 8) | (ord("F") << 16) | (ord("T") << 24)
	PROTOCOL_VER = 0x10000
	def __init__(self, direction):
		self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		self.direction = direction
		self.rate = 0
		self.cont = False
		self.force = False
		self.out = None
		self.ipv6 = False

	def connect(self, host, port):
		"""Connect to host:port"""
		print "Connecting to " + host + ":" + str(port)
		tries = 1
		while True:
			try:
				self.sock.connect((host, port))
				self.mode = "act"
				return True
			except socket.error, e:
				if tries < 3: time.sleep(1)
				else:
					print "Can't connect: " + e[1]
					return False
			tries += 1
	
	def listen(self, port):
		"""Listen for a connection on port"""
		print "Listening on " + str(port)
		servsock = self.sock
		servsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		try:
			if self.ipv6:
				servsock.bind(("::0", port))
			else:
				servsock.bind(("0.0.0.0", port))
			servsock.listen(1)
		except socket.error, e:
			print "Can't listen: " + e[1]
			return False
		self.sock, addr = servsock.accept()
		print "Got connection from " + addr[0] + ":" + str(addr[1])
		servsock.close()
		self.mode = "pasv"
		return True
	
	def send(self, data):
		"""Send data"""
		try:
			self.sock.sendall(data)
			return True
		except socket.error, e:
			print "Can't send: " + e[1]
			return False
	
	def recv(self, bytes):
		"""Receive 'bytes' bytes"""
		try:
			ret = ""
			while bytes > 0:
				data = self.sock.recv(bytes)
				if len(data) == 0:
					print "Connection lost"
					return False
				ret += data
				bytes -= len(data)
			return ret
		except socket.error, e:
			print "Can't receive: " + e[1]
			return False

	def sendPack(self, fmt, *data):
		"""Helperwrapper for sending structs"""
		return self.send(struct.pack("<" + fmt, *data))

	def recvPack(self, fmt):
		"""Helperwrapper for receiving structs"""
		raw = self.recv(struct.calcsize(fmt))
		if raw is False: return False
		data = struct.unpack("<" + fmt, raw)
		if len(data) == 1: return data[0]
		return data

	def actGreet(self):
		"""Active mode greeting"""
		if self.sendPack("iib", self.PROTOCOL_ID, self.PROTOCOL_VER, self.direction) is False: return False
		greetback = self.recvPack("i")
		if greetback is False or greetback != self.PROTOCOL_ID: return False
		return True

	def pasvGreet(self):
		"""Passive mode greeting"""
		greetback = self.recvPack("iib")
		if greetback is False: return False
		protid, protver, protdir = greetback
		if protid != self.PROTOCOL_ID or protver != self.PROTOCOL_VER or protdir == self.direction:
			return False
		return self.sendPack("i", self.PROTOCOL_ID)
	
	def greet(self):
		"""Handle greeting"""
		if self.mode == "act": return self.actGreet()
		else: return self.pasvGreet()

	def init(self, host, port):
		"""Make connection and greet"""
		if host is None:
			if not self.listen(port): return False
		else:
			if not self.connect(host, port): return False
		if not self.greet():
			print "Greeting failed"
			return False
		return True

	def setRate(self, rate):
		"""Set maximum transfer rate"""
		self.rate = rate
		if rate > 0: self.wanted_frame_time = 1024. / self.rate

	def limitRate(self):
		"""Sleep enough to get wanted rate"""
		frame_end = time.time()
		frame_time = frame_end - self.frame_start
		delay = self.wanted_frame_time - frame_time
		if delay > 0: time.sleep(delay)
		self.frame_start = frame_end + delay

	def close(self):
		self.sock.close()

	def setOpts(self, cont, force, out):
		self.cont = cont
		self.force = force
		self.out = out
	
	def enableIPv6(self):
		self.ipv6 = True
		self.sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)


class SsftSender(SsftSock):
	def __init__(self):
		SsftSock.__init__(self, 0)
		
	def sendFile(self, filename, handle):
		"""Send filename, read the contents from handle"""
		print "Sending " + filename
		size = fileSize(handle)
		filename = basename(filename)
		if self.sendPack("h", len(filename)) is False: return False
		# [2:] strips off the byte order marker, \xff\xfe
		newname = filename.decode(self.ENCODING, "replace").encode("utf-16")[2:]
		if self.send(newname) is False or self.sendPack("Q", size) is False: return False
		bytes = self.recvPack("Q")
		if bytes is False: return False
		h = hashlib.md5()
		if bytes != 0:
			print "Calculating hash of the first " + formatSize(bytes)
			h = calcHash(handle, bytes)
			if h is None:
				print "Error reading file"
				return False
			print "Waiting for answer"
			if self.send(h.digest()) is False: return False
			cont = self.recvPack("b")
			if cont is False:
				print "Peer did not accept the file"
				return False
			if cont == 0:
				handle.seek(0)
				h = hashlib.md5()
				bytes = 0
		sendlen = size - bytes
		print "Transferring " + formatSize(sendlen)
		self.frame_start = time.time()
		last_update = self.frame_start
		transferred = handle.tell()
		bar = Statusbar(size, transferred)
		data = handle.read(1024)
		while data:
			h.update(data)
			h.update(data)
			if self.send(data) is False: return False
			transferred += len(data)
			now = time.time()
			if now - last_update > 1:
				bar.update(transferred)
				last_update = now
			data = handle.read(1024)
			if self.rate > 0: self.limitRate()
		bar.update(transferred)
		if self.send(h.digest()) is False: return False
		handle.close()
		self.sock.close()
		return True

class SsftReceiver(SsftSock):
	def __init__(self):
		SsftSock.__init__(self, 1)
		
	def receiveFile(self):
		"""Receive file"""
		flen = self.recvPack("h")
		if flen is False: return False
		filename = self.recv(flen * 2)
		if filename is False: return False
		filename = filename.decode("utf-16").encode(self.ENCODING, "replace")
		size = self.recvPack("Q")
		if size is False: return False
		print "Peer offering file " + filename + " (" + formatSize(size) + ")"
		if self.out:
			filename = self.out
			print "Using file " + filename
		elif not self.force:
			try:
				newname = raw_input("Other filename? (empty = no): ")
				if newname != "": filename = newname
			except EOFError:
				pass
			except KeyboardInterrupt:
				print
				return None
		try:
			handle = file(filename, "ab+")
			handle.seek(0)
		except IOError, s:
			print "Can't open! " + s.strerror
			return False
		localsize = fileSize(handle)
		over = False
		if not self.force and localsize > size:
			if raw_input("Files do not match. Overwrite? ('y'?) ") == "y":
				over = True
			else:
				return False
		if self.sendPack("Q", localsize) is False: return False
		mhash = hashlib.md5()
		if localsize > 0:
			print "Calculating hash"
			mhash = calcHash(handle, localsize)
			if mhash is None: return False
			remote_hash = self.recv(16)
			if remote_hash is False: return False
			if self.force and not self.cont:
				over = True
			elif mhash.digest() != remote_hash:
				if raw_input("Files do not match. Overwrite? ('y'?) ") == "y":
					over = True
				else:
					return None
			elif localsize == size:
				if raw_input("Files match. Overwrite anyway? ('y'?) ") == "y":
					over = True
				else:
					return None
			else:
				if self.cont:
					over = False
				elif raw_input("Resume transfer? ('y'?) ") != "y":
					over = True
			if over:
				handle.seek(0)
				handle.truncate()
				mhash = hashlib.md5()
				cont = 0
				localsize = 0
			else:
				cont = 1
			if self.sendPack("b", cont) is False: return False
		print "Starting transfer"
		self.frame_start = time.time()
		last_update = self.frame_start
		bar = Statusbar(size, localsize)
		while localsize != size:
			bytes = size - localsize
			if bytes > 1024: bytes = 1024
			data = self.recv(bytes)
			if data is False: return False
			mhash.update(data)
			mhash.update(data)
			handle.write(data)
			localsize += bytes
			now = time.time()
			if now - last_update > 1:
				bar.update(localsize)
				last_update = now
			if self.rate > 0: self.limitRate()
		print
		bar.update(localsize)
		handle.close()
		remote_hash = self.recv(16)
		if remote_hash is False: return False
		self.sock.close()
		if mhash.digest() != remote_hash:
			print "File hashes do not match. There was maybe an error (or, more likely, a bug).\nYou might want to retry transfer to be sure about this."
		return True

class Statusbar:
	"""Wget-style file progress statusbar"""
	def __init__(self, totalsize, bytes):
		self.size = totalsize
		self.bytes_size = len(commify(totalsize))
		self.spd_size = len("xxxx.xxKB/s")
		self.eta_size = len("xx:yy:zz")
		self.cols = 80
		self.bar_size = self.cols - len("100% [] ") - self.bytes_size - 1 - self.spd_size - len(" ETA ") - self.eta_size - 1 - 1
		self.last_update = 0
		self.last_size = 0
		self.update(bytes)

	def update(self, bytes):
		"""Calculate things and print the bar"""
		now = time.time()
		if self.last_update == 0:
			self.last_update = now
			self.last_size = bytes
			return
		ratio = 1.0 * bytes / self.size
		percent = int(100 * ratio)
		printsize = commify(bytes)
		printsize += (" " * (self.bytes_size - len(printsize)))
		bar = "=" * int(ratio * self.bar_size + .5) + ">" + " " * int((1 - ratio) * self.bar_size + .5)
		if (now - self.last_update) != 0:
			bps = (bytes - self.last_size) / (now - self.last_update)
		else:
			bps = 0
		spd = formatSize(bps) + "/s"
		spd += " " * (self.spd_size - len(spd))
		if bps == 0:
			eta = ""
		else:
			eta = int((self.size - bytes) / bps)
			h = eta / 3600
			if h > 99:
				d = h / 24
				h %= 24
				eta = "%dd %dh"
			else:
				eta %= 3600
				m = eta / 60
				s = eta % 60
				eta = h and "%d:%02d:%02d" % (h, m, s) or "%02d:%02d" % (m, s)
		eta += " " * (self.eta_size - len(eta))
		print "\r%3d" % percent + "% [" + bar + "] " + printsize + " " + spd + " ETA " + eta,
		sys.stdout.flush()
		self.last_update = now
		self.last_size = bytes

def calcHash(f, num):
	"""Calculate hash of next num bytes in file f"""
	m = hashlib.md5()
	bar = Statusbar(num, 0)
	done = 0
	lasttime = time.time()
	while num > 4096:
		data = f.read(4096)
		if not data: return None
		m.update(data)
		done += 4096
		now = time.time()
		if now - lasttime > 1:
			bar.update(done)
			lasttime = now
		num -= 4096
	if num > 0:
		data = f.read(num)
		if not data: return None
		m.update(data)
		bar.update(done + num)
	return m

def basename(s):
	try:
		pos = s.rindex("/")
	except ValueError:
		try:
			pos = s.rindex("\\")
		except ValueError:
			return s
	return s[pos + 1:]

def fileSize(f):
	"""Get file size"""
	f.seek(0, 2)
	size = f.tell()
	f.seek(0)
	return size

def formatSize(n):
	"""Format a number to a human-readable form using suffixes, eg. 1024 -> 1K"""
	prefixes = ["", "K", "M", "G", "T", "P"]
	index = 0
	while n >= 1024 and index < 5:
		n /= 1024.
		index += 1
	n = round(n, 2)
	if n == int(n): n = int(n)
	return str(n) + prefixes[index] + "B"

def commify(val):
	"""Simplify reading using commas, eg. 1234 => 1,234"""
	val = str(val)
	n = len(val) - 1
	ret = ""
	for c in val:
		ret += c
		if (n % 3) == 0 and n != 0: ret += ","
		n -= 1
	return ret

def usage():
	"""Print help list"""
	name = sys.argv[0]
	print "Usage: %s [OPTIONS] [FILES]" % name
	print "Transfer file using 'ssft' protocol, see http://www.whitsoftdev.com/ssft/"
	print
	print " [OPTIONS]:"
	print "  -p, --port: tcp port to use"
	print "  -h, --host: host to connect to (omit to act as a server)"
	print "  -r, --rate: limit maximum transfer rate (in bytes per second with no suffix,"
	print "              KB/s with k etc.) Decimal values are allowed."
	print "  -o, --out: file(s) to write in when receiving; give repeatedly for many files"
	print "  -f, --force: use the name given and overwrite without prompting"
	print "  -c, --continue: resume downloading without prompting"
	print "  -l, --loop: receive files in a loop"
	print "  -6, --ipv6: send or receive on an ipv6 address"
	print " [FILES]: files to send; if none given, trying to receive."
	print 
	print "Examples:"
	print "  send files foo and bar to my.site.tld:1234:"
	print "    %s -p 1234 -h my.site.tld foo bar" % name
	print "  connect to my.site.tld:1234, receive two files and put them to foo and bar:"
	print "    %s -p 1234 -h my.site.tld -o foo -o bar" % name
	print "  listen on port 1234, then send file foo no faster than 1.5MB/s:"
	print "    %s -p 1234 -r 1.5m foo" % name
	print "  listen and receive whatever comes without prompting, continue when possible:"
	print "    %s -lfcp 1234" % name

def recv(ssft, host, port):
	"""Handle receiving"""
	if not ssft.init(host, port): return 1
	ret = ssft.receiveFile()
	if ret is True:
		print "Transfer complete."
	elif ret is False:
		print "Error in transfer."
	return 0

def send(ssft, host, port, name):
	"""Handle sending"""
	handle = None
	try:
		handle = file(name, "rb")
	except IOError:
		print "Can't open " + name
		return 0
	if not ssft.init(host, port): return 1
	ssft.sendFile(name, handle)
	print "Transfer complete."
	return 0

def main(args):
	"""Parse command line arguments and do the transfer"""
	try:
		opts, args = getopt.getopt(sys.argv[1:], "h:p:r:o:cfl6", ["host=", "port=", "rate=", "out=", "continue", "force", "loop", "ipv6"])
	except getopt.GetoptError, e:
		usage()
		return 1
	if len(opts) == 0 and len(args) == 0:
		usage()
		return 0
	host = None
	port = None
	rate = None
	cont = False
	force = False
	loop = False
	ipv6 = False
	out = []
	for opt, arg in opts:
		if opt in ("-h", "--host"):
			host = arg
		elif opt in ("-p", "--port"):
			try:
				port = int(arg)
			except ValueError:
				print "Bad port " + arg
				return 1
			if port < 0 or port > 65535:
				print "Valid port range 1..65535"
				return 1
		elif opt in ("-r", "--rate"):
			suffix = arg[-1:].lower()
			try:
				if suffix >= "0" and suffix <= "9":
					rate = float(arg)
				elif suffix in ("k", "m", "g"):
					rate = {"k":1024, "m":1048576, "g":1073741824}[suffix] * float(arg[:-1])
				else:
					print "Bad suffix " + suffix + " in rate"
					return 1
			except ValueError:
				print "Bad rate " + arg
				return 1
			if rate < 0:
				print "Bad rate " + arg
				return 1
		elif opt in ("-c", "--continue"):
			cont = True
		elif opt in ("-f", "--force"):
			force = True
		elif opt in ("-6", "--ipv6"):
			ipv6 = True
		elif opt in ("-l", "--loop"):
			loop = True
		elif opt in ("-o", "--out"):
			out.append(arg)
	if port is None:
		print "Error! No port given."
		usage()
		return 1
	if len(args) == 0:
		idx = 0
		while True:
			ssft = SsftReceiver()
			if rate: ssft.setRate(rate)
			if ipv6: ssft.enableIPv6()
			outfile = idx < len(out) and out[idx] or None
			idx += 1
			ssft.setOpts(cont, force, outfile)
			try:
				ret = recv(ssft, host, port)
				if not loop and idx >= len(out): return ret
				elif ret: time.sleep(1)
			except KeyboardInterrupt:
				print
				ssft.close()
				return 0
	else:
		for name in args:
			try:
				ssft = SsftSender()
				if rate: ssft.setRate(rate)
				if ipv6: ssft.enableIPv6()
				ret = send(ssft, host, port, name)
				if ret > 0: return ret
			except KeyboardInterrupt:
				ssft.close()
				return 0
			time.sleep(1)

if __name__ == "__main__":
	sys.exit(main(sys.argv[1:]))