/*
  JSmooth: a VM wrapper toolkit for Windows
  Copyright (C) 2003 Rodrigo Reyes <reyes@charabia.net>
 
  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., 675 Mass Ave, Cambridge, MA 02139, USA.
 
 */

/*
 * PEResourceDirectory.java
 *
 * Created on 2 août 2003, 01:28
 */

package net.charabia.jsmoothgen.pe;

import java.util.*;
import java.io.*;
import java.nio.*;
import java.nio.channels.*;

/**
 *
 * @author  Rodrigo
 */
public class PEResourceDirectory
{
	
    /*
      typedef struct IMAGE_RESOURCE_DIRECTORY {
      uint32_t   Characteristics;
      uint32_t   TimeDateStamp;
      uint16_t    MajorVersion;
      uint16_t    MinorVersion;
      uint16_t    NumberOfNamedEntries;
      uint16_t    NumberOfIdEntries;
      }
     */
	
	public class DataEntry
	{
		long   OffsetToData;   // To update at each change
		long   Size;
		long   CodePage;   // never changed
		long   Reserved;   // never changed
		ByteBuffer Data;
		
		public DataEntry(ByteBuffer data)
		{
			this.Data = data;
			this.Size = data.capacity();
		}
		
		public DataEntry(FileChannel chan, long offset) throws IOException
		{
			long orgpos = chan.position();
			chan.position(PEResourceDirectory.this.m_offsetBase + offset);
			ByteBuffer buf = ByteBuffer.allocate(16);
			buf.order(ByteOrder.LITTLE_ENDIAN);
			chan.read(buf);
			buf.position(0);
			
			OffsetToData = buf.getInt();
			Size = buf.getInt();
			CodePage = buf.getInt();
			Reserved = buf.getInt();
			
			long datapos = PEResourceDirectory.this.m_master.PointerToRawData + (OffsetToData - PEResourceDirectory.this.m_master.VirtualAddress);
			Data = ByteBuffer.allocate((int)Size);
			Data.order(ByteOrder.LITTLE_ENDIAN);
			chan.position(datapos);
			chan.read(Data);
			Data.position(0);
			
			chan.position(orgpos);
		}
		
		public int diskSize()
		{
			int size = 16 + (int)this.Size;
			if ((size%4)>0)
				size += 4-(size%4);
			return size;
		}
		
		public void dump(PrintStream out, int level)
		{
			indent(level, out); out.println("OffsetToData="+OffsetToData);
			indent(level, out); out.println("Size="+Size);
			indent(level, out); out.println("CodePage="+CodePage);
			indent(level, out); out.println("Reserved="+Reserved);
			indent(level, out); out.print("Data={ ");
			for (int i=0; i<this.Data.capacity(); i++)
			{
				out.print("" + Integer.toHexString((byte)Data.get()) + ",");
			}
			out.println(" }");
			
		}
		
		private void indent(int level, PrintStream out)
		{
			for (int i=0; i<level; i++)
				out.print("    ");
		}
		
		public int buildBuffer(ByteBuffer buffer, long virtualBaseOffset, int dataOffset)
		{
			//			System.out.println("Building Data Entry buffer @ " + buffer.position() + " (" + dataOffset + ")");
			
			dataOffset = buffer.position() + 16;
			
			buffer.putInt((int)(dataOffset + virtualBaseOffset));
			buffer.putInt((int)Size);
			buffer.putInt((int)CodePage);
			buffer.putInt((int)Reserved);
			
			Data.position(0);
			buffer.put(Data);
			
			dataOffset += Size;
			if ((dataOffset%4)>0)
				dataOffset += (4-(dataOffset%4));
			
			return dataOffset;
		}
		
		public void setData(ByteBuffer data)
		{
			Data = data;
			Size = data.capacity();
		}
	}
	
	public class ResourceEntry
	{
		int Id;
		String Name;
		
		ImageResourceDirectory Directory;
		DataEntry Data;
		
		public ResourceEntry(int id, DataEntry data)
		{
			this.Id = id;
			this.Data = data;
		}
		
		public ResourceEntry(String name, DataEntry data)
		{
			this.Name = name;
			this.Data = data;
		}
		
		public ResourceEntry(int id, ImageResourceDirectory dir)
		{
			this.Id = id;
			this.Directory = dir;
		}
		
		public ResourceEntry(String name, ImageResourceDirectory dir)
		{
			this.Name = name;
			this.Directory = dir;
		}
		
		public ResourceEntry(FileChannel chan) throws IOException
		{
			//			System.out.println("Resource Entry Offset: " + chan.position());
			ByteBuffer buf = ByteBuffer.allocate(8);
			buf.order(ByteOrder.LITTLE_ENDIAN);
			chan.read(buf);
			buf.position(0);
			long orgchanpos = chan.position();
			int val = buf.getInt();
			long offsetToData = buf.getInt();
			// 			System.out.println("Entry: Val=" + val);
			// 			System.out.println("       Off=" + offsetToData);
			
			if (val < 0)
			{
				val &= 0x7FFFFFFF;
				Name = extractStringAt(chan, val);
				Id = -1;
				//				System.out.println("    String at " + val + " = " + Name);
			} else
			{
				Id = val;
			}
			
			if (offsetToData < 0)
			{
				offsetToData &= 0x7FFFFFFF;
				long orgpos = chan.position();
				chan.position(PEResourceDirectory.this.m_offsetBase+offsetToData);
				Directory = new PEResourceDirectory.ImageResourceDirectory(chan);
				chan.position(orgpos);
			} else
			{
				Data = new DataEntry(chan, offsetToData);
			}
		}
		
		public String extractStringAt(FileChannel chan, int offset) throws IOException
		{
			long orgchanpos = chan.position();
			chan.position(PEResourceDirectory.this.m_offsetBase + offset);
			
			ByteBuffer sizebuf = ByteBuffer.allocate(2);
			sizebuf.order(ByteOrder.LITTLE_ENDIAN);
			chan.read(sizebuf);
			sizebuf.position(0);
			
			int size = sizebuf.getShort();
			ByteBuffer buffer = ByteBuffer.allocate(size*2);
			buffer.order(ByteOrder.LITTLE_ENDIAN);
			chan.read(buffer);
			buffer.position(0);
			
			StringBuffer buf = new StringBuffer(size);
			for (int i=0; i<size; i++)
			{
				int c = buffer.getShort();
				buf.append((char)c);
			}
			
			chan.position(orgchanpos);
			return buf.toString();
		}
		
		public int diskSize()
		{
			int size = 8;
			if (Name != null)
				size += (Name.length()*2) + 2;
			
			if (Directory != null)
				size += Directory.diskSize();
			else if (Data != null)
				size += Data.diskSize();
			
			if ((size%4)>0)
				size += 4-(size%4);
			return size;
		}
		
		public void dump(PrintStream out, int level)
		{
			indent(level, out);
			if (this.Name != null)
				out.println("Name=" + Name);
			else
				out.println("Id=#" + Id);
			
			indent(level, out);
			if (this.Directory != null)
			{
				out.println("ENTRY: DIRECTORY POINTER");
				this.Directory.dump(out, level + 1);
			}
			else
			{
				out.println("ENTRY: DATA ENTRY");
				Data.dump(out, level+1);
			}
		}
		
		private void indent(int level, PrintStream out)
		{
			for (int i=0; i<level; i++)
				out.print("    ");
		}
		
		public int buildBuffer(ByteBuffer buffer, long virtualBaseOffset, int dataOffset)
		{
			//			System.out.println("Building Resource Entry buffer  " + Name + "/" + Id + " @ " + buffer.position() + " (" + dataOffset + ")");
			if (Name != null)
			{
				buffer.putInt((int)(dataOffset | 0x80000000));
				
				int stringoffset = dataOffset;
				ByteBuffer strbuf = ByteBuffer.allocate(Name.length()*2 + 2);
				strbuf.order(ByteOrder.LITTLE_ENDIAN);
				
				strbuf.putShort((short)Name.length());
				for (int i=0; i<Name.length(); i++)
				{
					strbuf.putShort((short)Name.charAt(i));
				}
				strbuf.position(0);
				
				long oldpos = buffer.position();
				buffer.position(dataOffset);
				buffer.put(strbuf);
				dataOffset += Name.length()*2 + 2;
				if ((dataOffset%4) != 0)
					dataOffset += 4-(dataOffset%4);
				buffer.position((int)oldpos);
			} else
			{
				buffer.putInt(Id);
			}
			
			if (Directory != null)
			{
				buffer.putInt((int)(dataOffset | 0x80000000));
				
				int oldpos = buffer.position();
				buffer.position(dataOffset);
				int dirsize = Directory.buildBuffer(buffer, virtualBaseOffset);
				dataOffset = dirsize;
				buffer.position(oldpos);
				
			} else if (Data != null)
			{
				buffer.putInt(dataOffset);
				int oldpos = buffer.position();
				buffer.position(dataOffset);
				dataOffset = Data.buildBuffer(buffer, virtualBaseOffset, dataOffset);
				buffer.position(oldpos);
			} else
			{
				throw new RuntimeException("Directory and Data are both null!");
			}
			
			return dataOffset;
		}
	}
	
	public class ImageResourceDirectory
	{
		long Characteristics; // uint32_t
		long TimeDateStamp; // uint32_t
		int MajorVersion; // uint16_t
		int MinorVersion; // uint16_t
		int NumberOfNamedEntries; // uint16_t
		int NumberOfIdEntries; // uint16_t
		Vector NamedEntries = new Vector();
		Vector IdEntries = new Vector();
		
		public ImageResourceDirectory()
		{
		}
		
		public ImageResourceDirectory(FileChannel chan) throws IOException
		{
			ByteBuffer header = ByteBuffer.allocate(16);
			header.order(ByteOrder.LITTLE_ENDIAN);
			chan.read(header);
			header.position(0);
			
			Characteristics = header.getInt();
			TimeDateStamp = header.getInt();
			MajorVersion = header.getShort();
			MinorVersion = header.getShort();
			NumberOfNamedEntries = header.getShort();
			NumberOfIdEntries = header.getShort();
			
			for (int i=0; i<NumberOfNamedEntries; i++)
			{
				ResourceEntry re = new ResourceEntry(chan);
				NamedEntries.add(re);
			}
			for (int i=0; i<NumberOfIdEntries; i++)
			{
				ResourceEntry re = new ResourceEntry(chan);
				IdEntries.add(re);
			}
		}
		
		public void addNamedEntry(ResourceEntry entry)
		{
			this.NamedEntries.add(entry);
		}
		
		public void addIdEntry(ResourceEntry entry)
		{
			this.IdEntries.add(entry);
		}
		
		public void addEntry(ResourceEntry entry)
		{
			if (entry.Name != null)
				addNamedEntry(entry);
			else
				addIdEntry(entry);
		}
		
		public void dump(PrintStream out, int level)
		{
			indent(level, out); out.println("Directory: ");
			indent(level, out); out.println("Characteristics=" + this.Characteristics);
			indent(level, out); out.println("TimeDateStamp=" + this.TimeDateStamp);
			indent(level, out); out.println("MajorVersion=" + this.MajorVersion);
			indent(level, out); out.println("MinorVersion=" + this.MinorVersion);
			indent(level, out); out.println("NumberOfNamedEntries=" + this.NumberOfNamedEntries);
			indent(level, out); out.println("NumberOfIdEntries=" + this.NumberOfIdEntries);
			indent(level, out); out.println("Named Entries:");
			for (int i=0; i<NumberOfNamedEntries; i++)
			{
				ResourceEntry re = (ResourceEntry)NamedEntries.get(i);
				re.dump(out, level+1);
			}
			indent(level, out); out.println("Id Entries:");
			for (int i=0; i<NumberOfIdEntries; i++)
			{
				ResourceEntry re = (ResourceEntry)IdEntries.get(i);
				re.dump(out, level+1);
			}
		}
		
		private void indent(int level, PrintStream out)
		{
			for (int i=0; i<level; i++)
				out.print("    ");
		}
		
		public int diskSize()
		{
			int size = 16;
			for (int i=0; i<this.NamedEntries.size(); i++)
			{
				ResourceEntry re = (ResourceEntry)NamedEntries.get(i);
				size += re.diskSize();
			}
			for (int i=0; i<this.IdEntries.size(); i++)
			{
				ResourceEntry re = (ResourceEntry)IdEntries.get(i);
				size += re.diskSize();
			}
			
			if ((size%4)>0)
				size += 4-(size%4);
			
			return size;
		}
		
		public int buildBuffer(ByteBuffer buffer, long virtualBaseOffset)
		{
			//			System.out.println("Building Directory Entry buffer @ " + buffer.position());
			
			buffer.putInt((int)this.Characteristics);
			buffer.putInt((int)this.TimeDateStamp);
			buffer.putShort((short)this.MajorVersion);
			buffer.putShort((short)this.MinorVersion);
			buffer.putShort((short)this.NamedEntries.size());
			buffer.putShort((short)this.IdEntries.size());
			
			int dataOffset = buffer.position() + (NamedEntries.size()*8) + (IdEntries.size()*8);
			
			for (int i=0; i<this.NamedEntries.size(); i++)
			{
				ResourceEntry re = (ResourceEntry)this.NamedEntries.get(i);
				dataOffset = re.buildBuffer(buffer, virtualBaseOffset, dataOffset);
			}
			
			for (int i=0; i<this.IdEntries.size(); i++)
			{
				ResourceEntry re = (ResourceEntry)this.IdEntries.get(i);
				dataOffset = re.buildBuffer(buffer, virtualBaseOffset, dataOffset);
			}
			
			buffer.position(dataOffset);
			return dataOffset;
		}
		
		public ResourceEntry getResourceEntry(String name)
		{
			// If name == null, get the first entry in lexical 
			// order. If no entry in lexical order, choose the 
			// lowest integer id entry.
			if (name == null)
			{
				if (NamedEntries.size() > 0)
				{
					return (PEResourceDirectory.ResourceEntry)NamedEntries.get(0);
				}
				if (IdEntries.size() > 0)	
				{
					return (PEResourceDirectory.ResourceEntry)IdEntries.get(0);
				}
				return null;
			}
			
			if ((name.length()>0) && (name.charAt(0)=='#'))
			{
				try
				{
					String nb = name.substring(1);
					int i = Integer.parseInt(nb);
					return getResourceEntry(i);
				} catch (Exception exc)
				{
					exc.printStackTrace();
				}
			}
			
			for (Iterator i=this.NamedEntries.iterator(); i.hasNext(); )
			{
				ResourceEntry re = (ResourceEntry)i.next();
				if (name.equals(re.Name))
				{
					return re;
				}
			}
			return null;
		}
		
		public ResourceEntry getResourceEntry(int id)
		{
			for (Iterator i=this.IdEntries.iterator(); i.hasNext(); )
			{
				ResourceEntry re = (ResourceEntry)i.next();
				if (id == re.Id)
				{
					return re;
				}
			}
			return null;
		}
		
	}
	
	PESection m_master;
	PEFile m_file;
	long m_offsetBase;
	
	PEResourceDirectory.ImageResourceDirectory m_root;
	
	/** Creates a new instance of PEResourceDirectory */
	public PEResourceDirectory(PEFile file, PESection sect) throws IOException
	{
		m_master = sect;
		m_file = file;
		m_offsetBase = sect.PointerToRawData;
		init();
		
		//		System.out.println("--------\nTOTAL SIZE: " + m_root.diskSize());
		
		//		System.out.println("\n\n");
	}
	
	public void init() throws IOException
	{
	    ///		System.out.println("RESOURCE INIT");
	    //		System.out.println("   Offset: " + m_master.PointerToRawData);
		FileChannel chan = m_file.getChannel();
		chan.position(m_master.PointerToRawData);
		PEResourceDirectory.ImageResourceDirectory dir = new PEResourceDirectory.ImageResourceDirectory(chan);
		//		System.out.println("-----------------\nDUMP\n---------------");
		m_root = dir;
		
		//		dir.dump(System.out, 0);
	}
	
	public void dump(PrintStream out)
	{
		m_root.dump(out, 0);
	}
	
	public int size()
	{
		return m_root.diskSize();
	}
	
	public ByteBuffer buildResource(long virtualBaseOffset)
	{
	    //		System.out.println("BUILDING RESOURCE / VIRTUAL: " + virtualBaseOffset);
		int resourceSize = m_root.diskSize();
		ByteBuffer resbuf = ByteBuffer.allocate(resourceSize);
		resbuf.order(ByteOrder.LITTLE_ENDIAN);
		resbuf.position(0);
		m_root.buildBuffer(resbuf, virtualBaseOffset);
		return resbuf;
	}
	
	public PEResourceDirectory.ImageResourceDirectory getRoot()
	{
		return m_root;
	}	
	
	public boolean replaceResource(String catId, int resourceId, int langId, ByteBuffer data)
	{
		ResourceEntry catEntry = m_root.getResourceEntry(catId);
		if ((catEntry != null) && (catEntry.Directory != null))
		{
			ResourceEntry identEntry = catEntry.Directory.getResourceEntry(resourceId);
			if ((identEntry != null) && (identEntry.Directory != null))
			{
				ResourceEntry langEntry = identEntry.Directory.getResourceEntry(langId);
				if ((langEntry != null) && (langEntry.Data != null))
				{
					DataEntry dataslot = langEntry.Data;
					dataslot.setData(data);
					return true;
				}
			}
		}
		return false;
	}
	
	public void addNewResource(String catId, String resourceId, String languageId, ByteBuffer data)
	{
		DataEntry dataEntry = new DataEntry(data);
		ResourceEntry languageEntry = buildResourceEntry(languageId, dataEntry);
		ImageResourceDirectory languageDir = new ImageResourceDirectory();
		
		languageDir.TimeDateStamp = 0x3F2CCF64;
		languageDir.addEntry(languageEntry);
		
		ResourceEntry identEntry = buildResourceEntry(resourceId, languageDir);
		
		ImageResourceDirectory identDir = new ImageResourceDirectory();
		identDir.TimeDateStamp = 0x3F2CCF64;
		identDir.addEntry(identEntry);
		
		ResourceEntry catEntry = buildResourceEntry(catId, identDir);
		m_root.addEntry(catEntry);
	}
	
	public DataEntry getData(String catId, String resourceId, String langId)
	{
		ResourceEntry catEntry = m_root.getResourceEntry(catId);
		if ((catEntry != null) && (catEntry.Directory != null))
		{
			ResourceEntry identEntry = catEntry.Directory.getResourceEntry(resourceId);
			if ((identEntry != null) && (identEntry.Directory != null))
			{
				ResourceEntry langEntry = identEntry.Directory.getResourceEntry(langId);
				if ((langEntry != null) && (langEntry.Data != null))
				{
					DataEntry dataslot = langEntry.Data;
					return dataslot;
				}
			}
		}
		return null;
	}
	
	public ResourceEntry buildResourceEntry(String id, DataEntry data)
	{
		if ((id.length() > 1) && (id.charAt(0) == '#'))
		{
			int intid = Integer.parseInt(id.substring(1));
			return new ResourceEntry(intid, data);
		}
		
		return new ResourceEntry(id, data);
	}
	
	
	public ResourceEntry buildResourceEntry(String id, ImageResourceDirectory dir)
	{
		if ((id.length() > 1) && (id.charAt(0) == '#'))
		{
			int intid = Integer.parseInt(id.substring(1));
			return new ResourceEntry(intid, dir);
		}
		
		return new ResourceEntry(id, dir);
	}
	
}
