#!/usr/bin/python
# -*- mode: python; coding: utf-8 -*-
# jpegwav - extract .wav files from certain .jpeg files
# Copyright © 2005 Göran Weinholt <goran@weinholt.se>
#
# 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 2 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 St, Fifth Floor, Boston, MA  02110-1301 USA

# I wrote this to extract "audio clips" put in pictures taken by my
# HP Photosmart R817 digital camera. There's no guarantee that it'll
# work for other camera models as well, but it's always possible.
# There are some hard-coded values in JpegToWav.handle_APP2() that
# might mess things up with other cameras.

# An idea for the future: allow for removal of the audio clip from
# the pictures. Trivial to implement.

import sys, optparse, struct

class JpegReader:
	"""This is a JPEG reader class. It calls self.handle_XXX(marker, length)
	for each block of the JPEG it finds."""
	markers = { 0x01: 'TEM', 0xc0: 'SOF0', 0xc1: 'SOF1',
				0xc2: 'SOF2', 0xc3: 'SOF3', 0xc4: 'DHT',
				0xc5: 'SOF5', 0xc6: 'SOF6', 0xc7: 'SOF7',
				0xc8: 'JPG', 0xc9: 'SOF9', 0xca: 'SOF10',
				0xcb: 'SOF11', 0xcc: 'DAC', 0xcd: 'SOF13',
				0xce: 'SOF14', 0xcf: 'SOF15', 0xd0: 'RST0',
				0xd1: 'RST1', 0xd2: 'RST2', 0xd3: 'RST3',
				0xd4: 'RST4', 0xd5: 'RST5', 0xd6: 'RST6',
				0xd7: 'RST7', 0xd8: 'SOI', 0xd9: 'EOI',
				0xda: 'SOS', 0xdb: 'DQT', 0xdc: 'DNL',
				0xdd: 'DRI', 0xde: 'DHP', 0xdf: 'EXP',
				0xe0: 'APP0', 0xe1: 'APP1', 0xe2: 'APP2',
				0xe3: 'APP3', 0xe4: 'APP4', 0xe5: 'APP5',
				0xe6: 'APP6', 0xe7: 'APP7', 0xe8: 'APP8',
				0xe9: 'APP9', 0xea: 'APP10', 0xeb: 'APP11',
				0xec: 'APP12', 0xed: 'APP13', 0xee: 'APP14',
				0xef: 'APP15', 0xf0: 'JPG0', 0xfd: 'JPG13',
				0xfe: 'COM' }
	def __init__(self, filename):
		self.jpeg = open(filename)
		self.first = True
		for m in self.markers:
			setattr(self, self.markers[m], m)
		# Check if it's a JPEG
		if self.readbyte() != 255 or self.readbyte() != self.SOI:
			raise IOError("not a JPEG file")

	def readbyte(self):
		return struct.unpack(">B", self.jpeg.read(1))[0]
	def readword(self):
		return struct.unpack(">H", self.jpeg.read(2))[0]
	def readblock(self, length):
		return self.jpeg.read(length)
	def close(self):
		self.jpeg.close()

	def _handle_marker(self, marker, length):
		try: code = self.markers[marker]
		except KeyError: code = hex(marker)
		method = "handle_" + code
		#print method
		if hasattr(self, method):
			getattr(self, method)(marker, length)

	def handle_SOS(self, marker, length): # Start of Scan (headers end)
		self.close()
	def handle_EOI(self, marker, length): # End of Image
		self.close()

	def parse(self):
		while not self.jpeg.closed:
			while True:					# skip padding
				padding = self.readbyte()
				if padding == 255: break
			while True:					# really get marker
				marker = self.readbyte()
				if marker != 255: break
			length = self.readword()
			pos = self.jpeg.tell()
			if length < 2:
				raise IOError("invalid JPEG header length")
			length -= 2
			self._handle_marker(marker, length)	# handle the marker
			if not self.jpeg.closed:
				self.jpeg.seek(pos + length) # seek past this block

class JpegToWav(JpegReader):
	"""This class reads a JPEG file and, if it finds one, writes
	a .wav file to the given output file."""
	def __init__(self, input, output):
		JpegReader.__init__(self, input)
		self.output = output
		self.outfile = None
		self.block = 0

	def write_output(self, data):
		if self.outfile == None:
			self.outfile = open(self.output, 'w')
		self.outfile.write(data)

	def handle_APP2(self, marker, length):
		data = self.readblock(length)
		# Check if this is FlashPix data
		if data[:5] != 'FPXR\000': return
		version = struct.unpack('>B', data[5])[0]
		if version != 0:
			print "unknown FPXR version:", version
			return
		datatype = struct.unpack('>B', data[6])[0]
		self.block += 1
		
		#if datatype == 1:
		#	# I've tried to parse the contents list according to the
		#	# EXIF specs and there seems to be something non-standard
		#	# about the contents listings in my camera...
		#	print "Contents list found (ignoring)"
		#elif datatype == 2:
		#	print "Stream data"

		# Ugly workaround for not being able to parse the contents list,
		# I suppose.
		if self.block == 4:
			print "WAV header probably found, writing..." 
			self.write_output(data[41:])
		elif self.block > 4:
			print "Writing WAV data..."
			self.write_output(data[13:])

def main():
	usage = "usage: %prog input.jpeg output.wav"
	usage += "\nExtract .WAV files from certain digital camera .JPEG files."
	parser = optparse.OptionParser(usage, version="%prog 1.0")
	(opts, args) = parser.parse_args()

	if len(args) < 2:
		parser.error("must specify input and output files")

	jpeg = JpegToWav(args[0], args[1])
	jpeg.parse()

if __name__ == "__main__":
	main()