import os, hashlib, struct, subprocess, fnmatch, shutil, urllib, array
import wx
import png

from Crypto.Cipher import AES
from Struct import Struct

from common import *


class WOD: #WiiOpticalDisc
	class fsentry:
		name = ""
		parent = None

		def __init__(self, name, parent):
			self.name = ""
			if(parent != None):
				self.parent = parent
		def path(self):
			return parent.path() + "/" + name

	class fsdir(fsentry):
		def __init__(self, name, parent):
			fsentry.__init__(self, name, parent)

	class fsfile(fsentry):
		size = 0
		offset = 0

	class discHeader(Struct):
		__endian__ = Struct.BE
		def __format__(self):
			self.discId = Struct.string(1)
			self.gameCode = Struct.string(2)
			self.region = Struct.string(1)
			self.makerCode = Struct.uint16
			self.h = Struct.uint8
			self.version = Struct.uint8
			self.audioStreaming = Struct.uint8
			self.streamingBufSize = Struct.uint8
			self.unused = Struct.uint8[14]
			self.magic = Struct.uint32
			self.title = Struct.string(64)
			self.hashVerify = Struct.uint8
			self.h3verify = Struct.uint8

	def __str__(self):
		ret = ''
		ret += '%s [%s%s%s]\n' % (self.discHdr.title, self.discHdr.discId, self.discHdr.gameCode, self.discHdr.region)
		if self.discHdr.region == 'P':
			ret += 'Region : PAL\n'
		elif self.discHdr.region == 'E':
			ret += 'Region : NTSC\n'
		elif self.discHdr.region == 'J':
			ret += 'Region : JAPANESE\n'
		ret += 'Version 0x%x Maker %i Audio streaming %x\n' % (self.discHdr.version, self.discHdr.makerCode, self.discHdr.audioStreaming)
		ret += 'Hash verify flag 0x%x H3 verify flag : 0x%x\n' % (self.discHdr.hashVerify, self.discHdr.h3verify)
		ret += 'Found %i partitions (table at 0x%x)\n' % (self.partitionCount, self.partsTableOffset)
		ret += 'Found %i channels (table at 0x%x)\n' % (self.channelsCount, self.chansTableOffset)
		ret += 'Partition %i opened (type 0x%x) at 0x%x)\n' % (self.partitionOpen, self.partitionType, self.partitionOffset)
		ret += 'Tmd at 0x%x\n' % self.tmdOffset
		ret += 'main.dol at 0x%x fst at 0x%x (%ib)' % (self.dolOffset, self.fstSize, self.fstOffset)

		return ret

	def __init__(self, f):
		self.f = f
		self.fp = open(f, 'rb')

		self.discHdr = self.discHeader().unpack(self.fp.read(0x400))
		if self.discHdr.magic != 0x5D1C9EA3:
			raise Exception('Wrong disc magic')

		self.fp.seek(0x40000)

		self.partitionCount = 1 + struct.unpack(">I", self.fp.read(4))[0]
		self.partsTableOffset = struct.unpack(">I", self.fp.read(4))[0] << 2

		self.channelsCount = struct.unpack(">I", self.fp.read(4))[0]
		self.chansTableOffset = struct.unpack(">I", self.fp.read(4))[0] << 2

		self.partitionOpen = -1
		self.partitionOffset = -1
		self.partitionType = -1

	def openPartition(self, index):
		if index > self.partitionCount:
			raise ValueError('Partition index too big')

		self.partitionOpen = index

		self.partitionOffset = self.partsTableOffset + (8 * self.partitionOpen)

		self.fp.seek(self.partsTableOffset + (8 * self.partitionOpen))

		self.partitionOffset = struct.unpack(">I", self.fp.read(4))[0] << 2
		self.partitionType = struct.unpack(">I", self.fp.read(4))[0]

		self.fp.seek(self.partitionOffset + 0x2A4)

		self.tmdSize = struct.unpack(">I", self.fp.read(4))[0]
		self.tmdOffset = struct.unpack(">I", self.fp.read(4))[0] >> 2

		self.certsSize = struct.unpack(">I", self.fp.read(4))[0]
		self.certsOffset = struct.unpack(">I", self.fp.read(4))[0] >> 2

		self.H3TableOffset = struct.unpack(">I", self.fp.read(4))[0] >> 2

		self.dataOffset = struct.unpack(">I", self.fp.read(4))[0] >> 2
		self.dataSize = struct.unpack(">I", self.fp.read(4))[0] >> 2

		self.fp.seek(self.partitionOffset)
		open('partHeader.bin', 'w+b').write(self.fp.read(0x440))

		self.fp.seek(self.partitionOffset + 0x420)

		self.dolOffset = 4 * struct.unpack(">I", self.fp.read(4))[0]

		self.fstOffset = 4 * struct.unpack(">I", self.fp.read(4))[0]
		self.fstSize = 4 * struct.unpack(">I", self.fp.read(4))[0]

	def getIsoBootmode(self):
		if self.discHdr.discId == 'R' or self.discHdr.discId == '_':
			return 2
		elif self.discHdr.discId == '0':
			return 1

	def getOpenedPartition(self):
		return self.partitionOpen

	def getOpenedPartitionOffset(self):
		return self.partitionOffset

	def getOpenedPartitionType(self):
		return self.partitionType

	def getPartitionsCount(self):
		return self.partitionCount

	def getChannelsCount(self):
		return self.channelsCount

	def getPartitionCerts(self):
		self.fp.seek(self.partitionOffset + self.certsOffset)
		return self.fp.read(self.certsSize)

	def getPartitionH3Table(self):
		self.fp.seek(self.partitionOffset + self.H3TableOffset)
		return self.fp.read(0x18000)

	def getPartitionTmd(self):
		self.fp.seek(self.partitionOffset + self.tmdOffset)
		return self.fp.read(self.tmdSize)

	def getPartitionTik(self):
		self.fp.seek(self.partitionOffset)
		return self.fp.read(0x2A4)

	def getPartitionApploader(self):
		self.fp.seek(self.partitionOffset + 0x2440)
		print '%s' % hexdump(self.fp.read(32))
		self.fp.seek(self.partitionOffset + 0x2460)


	def extractPartition(self, index, fn = ""):

		if(fn == ""):
			fn = os.path.dirname(self.f) + "/" + os.path.basename(self.f).replace(".", "_") + "_out"
		try:
			origdir = os.getcwd()
			os.mkdir(fn)
		except:
			pass
		os.chdir(fn)

		self.fp.seek(0x18)
		if(struct.unpack(">I", self.fp.read(4))[0] != 0x5D1C9EA3):
			self.fp.seek(-4, 1)
			raise ValueError("Not a valid Wii Disc (GC not supported)! Magic: %08x" % struct.unpack(">I", self.fp.read(4))[0])

		self.fp.seek(partitionoffs)

		tikdata = self.fp.read(0x2A3)
		open("tik").write(tikdata)
		self.tik = Ticket("tik")
		self.titlekey = self.tik.getTitleKey()

		tmdsz = struct.unpack(">I", self.fp.read(4))[0]
		tmdoffs = struct.unpack(">I", self.fp.read(4))[0]

		certsz = struct.unpack(">I", self.fp.read(4))[0]
		certoffs = struct.unpack(">I", self.fp.read(4))[0]

		h3offs = struct.unpack(">I", self.fp.read(4))[0] << 2
		h3sz = 0x18000

		dataoffs = struct.unpack(">I", self.fp.read(4))[0] << 2
		datasz = struct.unpack(">I", self.fp.read(4))[0] << 2
		if(tmdoffs != self.fp.tell()):
			raise ValueError("TMD is in wrong place, something is fucked...wtf?")

		tmddata = self.fp.read(tmdsz)
		open("tmd").write(tmddata)

		self.tmd = TMD("tmd")


		print tmd.getIOSVersion()


		fst.seek(dataoffs)



		os.chdir("..")
	def _recurse(self, parent, names, recursion):
		if(recursion == 0):
			pass