osm2city-scripts/worldbuild-packager.py

#! /usr/bin/python3
# Copyright (C) 2018-2020 Merspieler, merspieler _at_ airmail.cc
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import os
import sys
import hashlib
import re
import socket
from subprocess import run, Popen, STDOUT
from time import sleep

host = socket.gethostname()
port = 12345
output = ""

argc = len(sys.argv)
i = 1
first = 1
while i < argc:
	if sys.argv[i] == "--port":
		i += 1
		port = sys.argv[i]
	elif sys.argv[i] == "--host":
		i += 1
		host = sys.argv[i]
	elif sys.argv[i] == "-o" or sys.argv[i] == "--output":
		i += 1
		output = sys.argv[i]
	elif sys.argv[i] == "-h" or sys.argv[i] == "--help":
		print("usage: worldbuild-packager.py [OPTIONS]")
		print("Retrives done tiles from the manager and packages them")
		print("")
		print("  -o, --output   Output directory. Path MUST be absolute")
		print("      --host     Manager host")
		print("      --port     Manager port")
		print("  -h, --help     Shows this help and exit")
		sys.exit(0)
	else:
		if first == 1:
			first = 0
			name = sys.argv[i]
		else:
			print("Unknown option " + sys.argv[i])
			sys.exit(1)
	i += 1

if output == "":
	print("ERROR: No output directory given")
	sys.exit(1)

def norm(num, length):
	num = str(abs(num))
	while len(num) < length:
		num = "0" + num
	return num

def send_status(name, status):
	try:
		sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		sock.connect((host, port))
		sock.send(("set " + name + " " + status).encode())
		sock.close()
	except IOError:
		print("Unable to send status. Aborting...")
		sys.exit(1)

def get_job(action):
	try:
		sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		sock.connect((host, port))
		sock.send(("get " + action).encode())
		msg = sock.recv(128)
		sock.close()
		msg = msg.decode()
		match = re.match(r"n-pole|s-pole|[ew]\d{3}[ns]\d{2}|None", msg)
		if match != None:
			ret = match.group(0)
			if ret == "None":
				print("No job got asigned. Trying again in one hour.")
				sleep(3600)
				ret = get_job(action)
		else:
			print("Recived invalid job. Retrying in 10 seconds...")
			sleep(10)
			ret = get_job(action)
		return ret
	except IOError:
		print("Unable to get job. Retrying in 10 seconds...")
		sleep(10)
		ret = get_job(action)
		return ret

def get_sha1(fname):
	if os.path.isfile(fname):
		sha1 = hashlib.sha1()
		with open(fname, 'rb') as f:
			while True:
				data = f.read(65536)
				if not data:
					break
				sha1.update(data)
		return sha1.hexdigest()
	else:
		return None


def read_index(fname):
	index = {}
	with open(fname) as f:
		for line in f:
			match = re.match("(.:.*):(.*)", line)
			if match != None:
				index[match.group(1)] = match.group(2)
			else:
				match = re.match("(version|time|path):(.*)", line)
				if match != None:
					index[match.group(1)] = match.group(2)
	return index

def write_index(fname, index):
	with open(fname, 'w') as f:
		for element in index:
			f.write(element + ":" + str(index[element]) + "\n")


def update_index(output, name, name_major):
	if os.path.isfile(output + "/" + name_major + "/.dirindex"):
		index = read_index(output + "/" + name_major + "/.dirindex")
	else:
		index = {}
		index["version"] = 1
		index["path"] = name_major

	sha = get_sha1(output + "/" + name_major + "/" + name + ".zip")
	if sha != None:
		index["z:" + name] = sha

		write_index(output + "/" + name_major + "/.dirindex", index)

		if os.path.isfile(output + "/.dirindex"):
			main_index = read_index(output + "/.dirindex")
		else:
			main_index = {}
			main_index["version"] = 1
			main_index["path"] = ""

		sha = get_sha1(output + "/" + name_major + "/.dirindex")
		if sha != None:
			main_index["z:" + name] = sha

		write_index(output + "/.dirindex", main_index)


running = True
while running:
	name = get_job("done")

	try:
		if name == "n-pole":
			pass
			# TODO
			send_status(name, "packaged")
		elif name == "s-pole":
			pass
			# TODO
			send_status(name, "packaged")
		else:
			match = re.match(r"([ew])(\d{3})([ns])(\d{2})", name)
			ew = match.group(1)
			ew_val = int(match.group(2))
			ns = match.group(3)
			ns_val = int(match.group(4))

			ew_val_major = int(ew_val / 10) * 10
			if ew == "w":
				if ew_val_major != ew_val:
					ew_val_major += 10

			ns_val_major = int(ns_val / 10) * 10
			if ns == "s":
				if ns_val_major != ns_val:
					ns_val_major += 10

			name_major = ew + norm(ew_val_major, 3) + ns + norm(ns_val_major, 2)

			run("mkdir -p " + output + "/" + name_major, shell=True)

			run("rm -f " + output + "/" + name_major + "/" + name + ".zip", shell=True)

			print("Packaging " + name + "...")

			cmd = "bash -c 'cd projects/worldbuild/scenery && zip -rq " + output + "/" + name_major + "/" + name + ".zip {Details,Pylons,Buildings,Roads}/" + name_major + "/" + name + "/*'"
			with open(os.devnull, 'w') as fp:
				package = Popen(cmd , shell=True, start_new_session=True, stdout=fp)

			package.wait()

			print("Packaging " + name + " done")

			update_index(output, name, name_major)

			send_status(name, "packaged")

	except KeyboardInterrupt:
		print("Graceful shutdown triggered. To force immedate stop, press Ctrl+C again")
		running = False
		try:
			package.wait()
			update_index(output, name, name_major)
			send_status(name, "packaged")
		except KeyboardInterrupt:
			# TODO doesn't work
			print("Forcing shutdown...")
			package.terminate()
			sleep(5)
			package.kill()
			sys.exit(0)