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

from common import norm, get_area_name_by_tile, api_get_job, api_send_status

host = socket.gethostname()
port = 12345
api = None
api_token = None
input_dir = ""
output = ""
tmp_dir = "/tmp/o2c-packager"

argc = len(sys.argv)
i = 1
first = 1
while i < argc:
	if sys.argv[i] == "-o" or sys.argv[i] == "--output":
		i += 1
		output = sys.argv[i]
	if sys.argv[i] == "-i" or sys.argv[i] == "--input":
		i += 1
		input_dir = sys.argv[i]
	elif sys.argv[i] == "-a" or sys.argv[i] == "--api":
		i += 1
		api = sys.argv[i]
	elif sys.argv[i] == "-t" or sys.argv[i] == "--api-token":
		i += 1
		api_token = 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("  -i  --input      Input directory. Path MUST be absolute")
		print("  -o, --output     Output directory. Path MUST be absolute")
		print("  -a, --api        Manager api url")
		print("  -t, --api-token  Manager api token")
		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)

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

if not api:
	print("ERROR: No api given")
	sys.exit(1)

if not api_token:
	print("ERROR: No api token given")
	sys.exit(1)

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("(version|time|path|d:.*):(.*)", line)
			if match != None:
				index[match.group(1)] = match.group(2)
			else:
				match = re.match("(.*:.*):(.*:.*)", 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_parent_index(base, path, changed_dir):
	if os.path.isfile(base + "/" + path + "/.dirindex"):
		index = read_index(base + "/" + path + "/.dirindex")
	else:
		index = {}
		index["version"] = 1
		if path != "":
			index["path"] = path
		else:
			index["path"] = ""

	sha = get_sha1(base + "/" + path + "/" + changed_dir + "/.dirindex")
	if sha != None:
		index["d:" + changed_dir] = sha

	write_index(base + "/" + path + "/.dirindex", index)


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

	sha = get_sha1(output + "/" + part + "/" + name_major + "/" + name + ".txz")
	if sha != None:
		index["t:" + name + ".txz"] = sha + ":" + str(os.path.getsize(output + "/" + part + "/" + name_major + "/" + name + ".txz"))

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

		update_parent_index(output, part, name_major)
		update_parent_index(output, "", part)

# Always make sure, your tmp dirs exists
for part in os.listdir(input_dir):
	run("mkdir -p " + tmp_dir + "/" + part, shell=True)

previous_area = ""
running = True
while running:
	try:
		tiles = "None"
		while tiles == "None":
			tiles = api_get_job("done", "packaging", api, api_token, none_exit=False, all_in_parent=1)
			if tiles == "None":
				print("No job got asigned. Trying again in one hour.")
				sleep(3600)

		area = get_area_name_by_tile(tiles[0])
		area_major = get_area_name_by_tile(tiles[0], major=True)

		print("Packaging " + area + "...")

		for part in os.listdir(input_dir):
			tmp_dir_part = tmp_dir + "/" + part

			# Check if we can reuse the already extracted structure
			if previous_area != area:
				run("rm -rf " + tmp_dir_part + "/" + area, shell=True)
				# If there is an archive, we need to extract it first and remove old data
				if os.path.exists(output + "/" + part + "/" + area_major + "/" + area + ".txz"):
					run("bash -c 'cd " + tmp_dir_part + " && tar -xf " + output + "/" + part + "/" + area_major + "/" + area + ".txz'", shell=True)
					for name in tiles:
						name = str(name)
						run("bash -c 'cd " + tmp_dir_part + " && rm -f " + area + "/*" + name + "*'", shell=True)
			run("mkdir -p " + tmp_dir_part + "/" + area, shell=True)

			# If there's new scenery copy it into the tmp dir
			for name in tiles:
				name = str(name)
				if os.path.exists(input_dir + "/" + part + "/" + area_major + "/" + area + "/" + name + ".stg"):
					run("cp " + input_dir + "/" + part + "/" + area_major + "/" + area + "/*" + name + "* " + tmp_dir_part + "/" + area + "/", shell=True)

			# If the scenery is not empty, package it
			if os.listdir(tmp_dir_part + "/" + area):
				run("mkdir -p " + output + "/" + part + "/" + area_major, shell=True)

				cmd = "bash -c 'cd " + tmp_dir_part + " && tar -cJf " + output + "/" + area + ".txz " + area + "/'"
				with open(os.devnull, 'w') as fp:
					package = Popen(cmd , shell=True, start_new_session=True, stdout=fp)
				package.wait()
				run("mv " + output + "/" + area + ".txz " + output + "/" + part + "/" + area_major + "/", shell=True)

				update_index(output, part, area, area_major)


		print("Packaging " + area + " done")
		previous_area = area

		for name in tiles:
			name = str(name)
			api_send_status(name, "packaged", api, api_token)

	except KeyboardInterrupt:
		print("Graceful shutdown triggered. To force immedate stop, press Ctrl+C again")
		running = False
		try:
			package.wait()
			run("mv " + output + "/" + area + ".txz " + output + "/" + part + "/" + area_major + "/", shell=True)
			update_index(output, area, area_major)
			api_send_status(name, "packaged", api, api_token)
		except KeyboardInterrupt:
			# TODO doesn't work
			print("Forcing shutdown...")
			package.terminate()
			sleep(5)
			package.kill()
			sys.exit(0)