dawn/updater/sophon/updater.py

#!/usr/bin/env python3

"""
	Functional description
	----------------------
	chunks
		zstd-compressed sections of files for patching
	manifest
		Provides information for all chunks
		name (= URL), patch offset and checksum
	getBuild
		JSON file that provides information about the available game and voiceover pack files
		Provides manifests and the base URL to download chunks from


	TODO
	----
	High priority
		New game installation
			Which URL?
		Where is 'deletefiles.txt' ?
		voiceover-packs
	Medium priority
		Make chunk names identical to the official launcher
	Low priority
		Parallelization for file downloads and patching

"""



import hashlib # md5
import json
import manifest_pb2 # generated
import os
import pathlib
import re # Regular Expressions
import shutil # rmtree
import sys
import time
import urllib.request as request # downloads
import zstandard # archive unpacking
from google.protobuf.json_format import MessageToJson

TEMPDIR: pathlib.Path
GAMEDIR: pathlib.Path
SCRIPTDIR = os.path.dirname(os.path.realpath(__file__))

def tempdir(*args):
	global TEMPDIR
	return os.path.join(TEMPDIR, *args)

def gamedir(*args):
	global GAMEDIR
	return os.path.join(GAMEDIR, *args)

def infolog(*args):
	print("INFO ", *args)

def warnlog(*args):
	print("WARN ", *args)


class SophonClient:
	installed_ver: None  # "major.minor.patch" or "new" for new installations
	rel_type: None       # os / cn
	gamedatadir: None    # absolute path to *_Data
	branch: None         # main / predownload
	branches_json = None # package_id, password, tag
	getBuild_json = None # json object of the entire getBuild json
	manifest: None       # manifest_pb2
	category_json: None  # json object of currently selected category


	def initialize(self):
		global GAMEDIR, TEMPDIR

		# CLI argument parsing
		args = sys.argv[2 :]

		new_install = ("--install" in args)
		if new_install:
			# Check path sanity
			if not os.path.isdir(GAMEDIR):
				os.mkdir(GAMEDIR)
			else:
				# must be empty (allow config.ini)
				assert len(list(GAMEDIR.glob("*"))) < 2, "The specified path is not empty"

		if not os.path.isdir(GAMEDIR):
			print("Script usage: python3 updater.py /path/to/game/dir")
			exit(1)

		self.branch = "predownload" if ("--predownload" in args) else "main"

		if not os.path.isdir(TEMPDIR):
			os.mkdir(TEMPDIR)

		# Autodetection
		if new_install:
			self.initialize_install()
		else:
			self.initialize_update()


	def initialize_install(self):
		assert False, "TODO"
		# Prompt to ask for the desired game version
		pass


	def initialize_update(self):
		"""
		Find out what kind of installation we need to update
		"""
		global TEMPDIR, GAMEDIR

		# Absolute path to the game data directory
		datadir = list(GAMEDIR.glob("*_Data"))
		assert len(datadir) == 1, "Cannot determine game data dir"
		self.gamedatadir = datadir[0]

		if os.path.isfile(gamedir("GenshinImpact.exe")):
			self.rel_type = "os"
		elif os.path.isfile(gamedir("YuanShen.exe")):
			if os.path.isfile(self.gamedatadir.join("Plugins/PCGameSDK.dll")):
				self.rel_type = "bb"
			else:
				self.rel_type = "cn"

		assert isinstance(self.rel_type, str), "Failed to detect release type"
		infolog(f"Release type: {self.rel_type}")

		# Retrieve the installed game version
		if True:
			fullname = gamedir(self.gamedatadir, "globalgamemanagers")
			assert os.path.isfile(fullname), "Game install is incomplete!"

			contents = open(fullname, "rb").read()
			ver = re.findall(b"\0(\d+\.\d+\.\d+)_\d+_\d+\0", contents)
			assert len(ver) == 1, "Broken script or corrupted game installation"

			self.installed_ver = ver[0].decode("utf-8")
			infolog(f"Installed game version: {self.installed_ver}")

		# Compare game version with what's contained in "config.ini"
		self.check_config_ini()


	def check_config_ini(self):
		"""
		Internal function
		"""
		fullname = gamedir("config.ini")
		version_config = None
		if not os.path.isfile(fullname):
			return

		contents = open(fullname, "r").read()
		ver = re.findall(r"game_version=(\d+\.\d+\.\d+)", contents)
		if len(ver) != 1:
			warnlog("config.ini is incomplete or corrupt")
			return

		infolog(f"config.ini: Game version {ver[0]}")
		if ver[0] != self.installed_ver:
			warnlog("config.ini and the actual installed game version differ!")


	def cleanup_temp(self):
		shutil.rmtree(TEMPDIR)


	def load_or_download_file(self, fname, url):
		"""
		fname: file name without path prefix
		url:   str or function ptr to retrieve the URL
		Returns: File handle
		"""
		fullname = tempdir(fname)
		do_download = True

		if os.path.isfile(fullname):
			# keep cached for 24 hours
			do_download = time.time() - os.path.getmtime(fullname) > (24 * 3600)

		if do_download:
			# Check whether the file is still up-to-date
			if callable(url):
				url = url()
			try:
				os.mkdir(TEMPDIR)
			except FileExistsError:
				pass

			request.urlretrieve(url, fullname)
			infolog(f"Got file '{fname}' (new)")
		else:
			infolog(f"Got file '{fname}' (cached)")

		return open(fullname, "rb")


	def retrieve_API_keys(self):
		"""
		Retrieves passkeys for authentication to download URLs
		Depends on "initialize_*".
		"""

		assert isinstance(self.rel_type, str), "Missing initialize"

		base_url = None
		tail = None

		if self.rel_type == "os":
			# Up-to-date as of 2024-06-15 (4.7.0)
			game_ids = "gopR6Cufr3"
			launcher_id = "VYTpXlbWo8"
			tail = f"game_ids[]={game_ids}&launcher_id={launcher_id}"
			base_url = "https://sg-hyp-api.hoy" + "overse.com/hyp/hyp-connect/api"

		assert self.rel_type == "os", "CN/BB yet not implemented" # TODO

		if not self.branches_json:
			# MANDATORY. JSON with package_id, password and tag(s)
			fh = self.load_or_download_file("getGameBranches.json", f"{base_url}/getGameBranches?{tail}")
			js = json.load(fh)
			assert js["retcode"] == 0, "Failed to retrieve API keys"
			# Array length corresponds to the amount of "game_ids" requested.
			self.branches_json = js["data"]["game_branches"][0][self.branch]
			assert self.branches_json != None, f"Cannot find API keys for branch '{self.branch}'"

		if False:
			# JSON with game paths for voiceover packs, logs, screenshots
			self.load_or_download_file("getGameConfigs.json", f"{base_url}/getGameConfigs?{tail}")

		if False:
			# JSON with SDK files (BiliBili ?)
			channel = 1
			sub_channel = 0
			self.load_or_download_file("getGameChannelSDKs.json",
					f"{base_url}/getGameChannelSDKs?channel={channel}&{tail}&sub_channel={sub_channel}")

		if False:
			# zip downloads (successor of the /resource JSON file)
			self.load_or_download_file("getGamePackages.json", f"{base_url}/getGamePackages?{tail}")


	def make_getBuild_url(self):
		"""
		Compose the URL for the main JSON file for chunk-based downloads (sophon)
		Returns: URL
		"""
		if not self.branches_json:
			self.retrieve_API_keys()


		url = None
		if self.rel_type == "os":
			url = "sg-public-api.ho" + "yoverse.com"
		elif self.rel_type == "cn":
			url = "api-takumi.mih" + "oyo.com"

		assert not (url is None), f"Unhandled release type {self.rel_type}"

		url = (
			"https://" + url + "/downloader/sophon_chunk/api/getBuild"
			+ "?branch=" + self.branches_json["branch"]
			+ "&package_id=" + self.branches_json["package_id"]
			+ "&password=" + self.branches_json["password"]
		)

		infolog("Created getBuild JSON URL")
		return url


	def load_getBuild_json(self):
		"""
		Loads the main JSON for manifest and chunk information
		"""

		fh = self.load_or_download_file("getBuild.json", self.make_getBuild_url)
		self.getBuild_json = json.load(fh)
		infolog("Loaded getBuild JSON")


	def load_manifest(self, cat_name):
		"""
		Loads the specified manifest protobuf
		cat_name: "game", "en-us", "zh-cn", "ja-jp", "ko-kr"
		"""

		if not self.getBuild_json:
			self.load_getBuild_json()

		jd = self.getBuild_json["data"]
		infolog(f"Server provides game version {jd['tag']}")

		# Find matching category
		category = None
		for jdm in jd["manifests"]:
			if jdm["matching_field"] == cat_name:
				category = jdm
				break

		assert not (category is None), f"Cannot find the specified field '{cat_name}'"
		infolog(f"Found category {cat_name}")
		self.category_json = category

		# Download and decompress manifest protobuf
		fname_raw = category["manifest"]["id"]
		url = category["manifest_download"]["url_prefix"] + "/" + category["manifest"]["id"]

		zfh = self.load_or_download_file(fname_raw + ".zstd", url)
		reader = zstandard.ZstdDecompressor().stream_reader(zfh)
		pb = manifest_pb2.Manifest()
		pb.ParseFromString(reader.read())
		nfiles = len(pb.files)
		infolog(f"Decompressed manifest protobuf ({nfiles} files)")

		json_fname = tempdir(fname_raw + ".json")
		if not os.path.isfile(json_fname):
			with open(json_fname, "w+") as jfh:
				json.dump(json.loads(MessageToJson(pb)), jfh)
			infolog("Exported protobuf to JSON file")

		self.manifest = pb


	def find_chunks_by_file_name(self, file_name):
		"""
		Searches a specific file name in the manifest
		Returns: FileInfo or None
		"""
		assert isinstance(file_name, str)
		for v in self.manifest.files:
			if v.filename == file_name:
				return v

		warnlog(f"Cannot find chunks for file: {file_name}")
		return None


	def download_and_patch_file(self, file_info):
		"""
		Downloads the chunks and patches a file
		file_info: FileInfo, one of the manifest.files[] objects
		"""
		if file_info.flags == 64:
			# Created as soon a file is put inside
			infolog(f"Skipping directory entry: {file_info.filename}")
			return
		assert (file_info.flags == 0), f"Unknown flags {file_info.flags} for '{file_info.filename}'"

		chunk_url_prefix = self.category_json["chunk_download"]["url_prefix"]

		infolog(
			"Chunk downloader:\n"
			f"\t File name:   {file_info.filename}\n"
			f"\t Chunk count: {len(file_info.chunks)}\n"
			f"\t File size:   {int(file_info.size / (1024 * 1024 / 10) + 0.5) / 10} MiB"
		)

		# Disallow writing to unpredictable locations
		assert (".." not in str(file_info.filename)), "Security alert"
		assert (str(file_info.filename)[0] != '/'), "Security alert"

		# Create directory structure and patch file into the temporary directory
		parent_dir = pathlib.Path(file_info.filename).parent
		os.makedirs(tempdir(parent_dir), exist_ok=True)

		tmp_file = tempdir(file_info.filename)
		shutil.copyfile(gamedir(file_info.filename), tmp_file)
		fh = open(tmp_file, "wb") # no truncate!

		# Download all chunks
		n = 0
		for chunk in file_info.chunks:
			n += 1
			cfname = tempdir(chunk.chunk_id)
			try:
				if os.path.getsize(cfname) == chunk.compressed_size:
					continue # TODO: do a proper hash check
			except FileNotFoundError:
				pass # Already downloaded

			request.urlretrieve(chunk_url_prefix + "/" + chunk.chunk_id, cfname)
			print(f"\t * Downloaded chunk {n} / {len(file_info.chunks)}")

		# Apply all chunks
		n = 0
		for chunk in file_info.chunks:
			n += 1
			cfname = tempdir(chunk.chunk_id)

			# Attempt to patch
			zfh = open(cfname, "rb")
			reader = zstandard.ZstdDecompressor().stream_reader(zfh)
			fh.seek(chunk.offset)
			fh.write(reader.read())

			print(f"\t * Patched {n} / {len(file_info.chunks)}")

		fh.close() # patch done

		# Verify file integrity
		if file_info.md5 == hashlib.md5(open(tmp_file, "rb").read()).hexdigest():
			print("\t * Patching succeeded. md5 matches.")
		else:
			print(f"\t * Hash mismatch after patching. File is corrupt: {file_info.filename}")
		print("")

# ------------------- Actions

# This script does currently not write anything to the game installation directory
GAMEDIR = pathlib.Path(sys.argv[1])
TEMPDIR = os.path.join(SCRIPTDIR, "tmp")

cli = SophonClient()
cli.initialize()
cli.retrieve_API_keys()

#cli.cleanup_temp()
cli.load_manifest("game")

v = cli.find_chunks_by_file_name("pkg_version")
if v:
	cli.download_and_patch_file(v)