#! /usr/bin/python
#
# Generates striped MD5 fingerprints for BF2 content checking.
#
# Author: Andreas Fredriksson
# Copyright (c)2005 Digital Illusions CE AB

import hashlib
import os
import os.path
from optparse import OptionParser

num_ordinals = 10
block_size = 16384

def ordinal_pass(files):
	ctx = []
	for x in xrange(0, num_ordinals):
		ctx.append(hashlib.md5())

	for file in files:
		sz = os.stat(file).st_size
		num_blocks = (sz / block_size) + 1

		print " Fingerprinting ", file		
		inf = open(file, 'rb')
		for b_i in xrange(0, num_blocks+1, num_ordinals):
			for ordinal in xrange(0, num_ordinals):
				data = inf.read(block_size)
				if len(data) > 0: 
					ctx[ordinal].update(data)
		inf.close()
	
		#Take size into account as well.
		for ordinal in xrange(0, num_ordinals):	
			ctx[ordinal].update(str(sz))
	
	return ctx

def run_fingerprint(files):
	ctx = ordinal_pass(files)
	tmp = []
	for o in xrange(0, num_ordinals):
		tmp.append('%d %s' % (o, ctx[o].hexdigest()))
	return tmp
	
def generate_md5(md5File, fileList, level = None):
	"""
	Compute MD5 over all files in filelist, and write that to the given outfile.

	If level is non-None, it is prepended to each output line, this format is used for levels.
	"""
	print "Building MD5 file \"%s\":" % md5File
	fingerprint = run_fingerprint(fileList)
	of = open(md5File, 'wt')
	for s in fingerprint:
		if not level is None:
			of.write(level)
			of.write(" ")
		of.write(s)
		of.write('\n')
	of.close()
	

def visitor(arcs, dirname, names):
	for fn_ in names:
		fn = fn_.lower()
		if fn == 'server.zip' or fn == 'client.zip':
			p = os.path.split(dirname)[1].lower()
			list = arcs.get(p, [])
			list.append(os.path.join(dirname, fn))
			arcs[p] = list
	return names

parser = OptionParser()

parser.add_option(
	"--levels",
	action="store_true",
	dest="levels",
	default=False,
	help="Just checksum levels only")

parser.add_option(
	"--root",
	action = "store",
	dest = "root",
	default = None,
	help = "The root bin/ directory of the game to checksum")

(options, args) = parser.parse_args()

try:
	mod_name = args[0]
except:
	print("\n****** Content Checking: Mod name missing - using mods/bf2 ******\n")
	mod_name = "bf2"

# New root directory specified?
if options.root != None:
	os.chdir(options.root)

mod_path = os.path.join("mods", mod_name)

files = []
files.append(os.path.join(mod_path, "ClientArchives.con"))
files.append(os.path.join(mod_path, "ServerArchives.con"))

std_archives = []
nite_archives = []
boost_archives = []
boost_nite_archives = []

if options.levels == False:
	print("\n****** Content Checking: Collecting " + mod_name + " common files to fingerprint... ******\n")
	
	for filename in files :
		of = open(filename,'rt')
		for line in of:
			list = line.split(' ')
			if len(list) > 1 :
				archive = list[1]
				archive = archive.replace('\\', os.sep)
				archive = archive.replace('/', os.sep)	  
				if archive[:4] != 'mods' :
					archive = os.path.join(mod_path, archive)
				archive = archive.lower()
				if archive.find( 'shaders' ) >= 0:
					print("\nContent Checking: Found SHADER updating bf2 archives list...\n")
					bf2archive = "mods\\bf2\\shaders_client.zip"
					std_archives += [bf2archive]		# Use bf2 shaders as day shaders
					nite_archives += [archive]			# Use mods shaders as nite shaders
					boost_archives += [bf2archive]
					boost_nite_archives += [archive]
				else:
					if archive.find( 'booster' ) >= 0:
						print("\nContent Checking: Found BOOSTER ignoring for standard & nite archives list...\n")
						boost_archives += [archive]			 # only add to boster lists
						boost_nite_archives += [archive]
					else:
						std_archives += [archive]
						nite_archives += [archive]
						boost_archives += [archive]
						boost_nite_archives += [archive]
		of.close()
	
	std_archives.sort()
	nite_archives.sort()
	boost_archives.sort()
	boost_nite_archives.sort()
	
	print("Day Archives")
	for archive in std_archives :
		print " ", archive
	
	print("Night Archives")
	for narchive in nite_archives :
		print " ", narchive
	
	print("Booster Archives")
	for barchive in boost_archives :
		print " ", barchive
	
	print("Booster+Night Archives")
	for bnarchive in boost_nite_archives :
		print " ", bnarchive
	
	print("\n****** Content Checking: Processing " + mod_name + " common archives for DAY... ******\n")
	generate_md5(os.path.join(mod_path, 'std_archive.md5'), std_archives)

	print("\n****** Content Checking: Processing " + mod_name + " common archives for NIGHT... ******\n")
	generate_md5(os.path.join(mod_path, 'std_archive_mod.md5'), nite_archives)

	print("\n****** Content Checking: Processing " + mod_name + " common archives for BOOSTER... ******\n")
	generate_md5(os.path.join(mod_path, 'bst_archive.md5'), boost_archives)

	print("\n****** Content Checking: Processing " + mod_name + " common archives for BOOSTER+NIGHT... ******\n")
	generate_md5(os.path.join(mod_path, 'bst_archive_mod.md5'), boost_nite_archives)

print("\n****** Content Checking: Collecting " + mod_name + " files to fingerprint... ******\n")

lvl_archives = {}
os.path.walk(os.path.join(mod_path, 'levels'), visitor, lvl_archives)

print("\n****** Content Checking: Found " + str(len(lvl_archives)) + " " + mod_name + " Levels ******\n")

print("\n****** Content Checking: Processing " + mod_name + " Level archives... ******\n")

for level, arcs in lvl_archives.iteritems():
	arcs.sort()
	generate_md5(os.path.join(mod_path, 'levels', level, 'archive.md5'), arcs, level = level)

print("\n****** Content Checking: " + mod_name + " Done ******\n")