file_info.nim

from parseopt import getopt, cmdArgument, cmdShortOption
from strutils import `%`
from algorithm import sorted
from streams import newFileStream, Stream, setPosition,
                    getPosition, close, write, readStr,
                    readChar, atEnd, newStringStream
import layout

type
  NumberedObject = object
    idx: int
    obj: ObjectInfo

  AlignedStreamObj = object of Stream
    shift: int
    oldstream: Stream
  AlignedStream = ref AlignedStreamObj

var filename = ""
var outfile = ""
var sourcefile = ""

for kind, key, val in getopt():
  case kind
  of cmdArgument:
    filename = key
  of cmdShortOption:
    case key
    of "o": outfile = val
    of "s": sourcefile = val
    else:
      discard
  else:
    discard


echo "Target: ", filename
let reader = newFileStream(filename, fmRead)
echo "Source: ", sourcefile
let sreader = newFileStream(sourcefile, fmRead)

proc asr(s:Stream,buffer:pointer,bufLen:int):int =
  let s = AlignedStream(s)
  var pos = s.oldStream.getPosition() - s.shift
  while (pos mod 4) != 0:
    pos += 1
  s.oldStream.setPosition(pos + s.shift)
  s.oldStream.readDataImpl(s.oldStream, buffer, bufLen)

proc asw(s:Stream,buffer:pointer,bufLen:int) =
  let s = AlignedStream(s)
  var c = '\x00'
  var pos = s.oldStream.getPosition() - s.shift
  while (pos mod 4) != 0:
    s.oldStream.writeDataImpl(s.oldStream, c.addr, 1)
    pos += 1
  s.oldStream.setPosition(pos + s.shift)
  s.oldStream.writeDataImpl(s.oldStream, buffer, bufLen)

proc get_aligned_stream(ss: Stream, val: int): AlignedStream =
  var a: AlignedStream = new(AlignedStream)
  a.shift = ss.getPosition()
  a.setPositionImpl = ss.setPositionImpl
  a.getPositionImpl = ss.getPositionImpl
  a.oldStream = ss
  a.readDataImpl = cast[type(a.readDataImpl)](asr)
  a.writeDataImpl = cast[type(a.writeDataImpl)](asw)
  a
  #result.readDataImpl = cast[type(s.readDataImpl)](treader)
  #result.writeDataImpl = cast[type(s.writeDataImpl)](twriter)

proc write_file(s: Stream, f: FileInfo) =
  let aligned = get_aligned_stream(s, 4)
  f.serialize(aligned)

proc get_file(source: Stream, begin: uint32, s: ObjectInfo): FileInfo =
  source.setPosition(int(s.byteStart + begin))
  let aligned = get_aligned_stream(source, 4)
  FileInfo.deserialize(aligned)

proc compare(a: ObjectInfo, b: ObjectInfo): int =
  cmp(a.byteStart, b.byteStart)

proc compare(a: NumberedObject, b:NumberedObject): int =
  cmp(a.obj.byteStart, b.obj.byteStart)

proc extract(source: Stream, begin: uint32, s: ObjectInfo) =
  let finfo = get_file(source, begin, s)
  echo "Filename: ", finfo.name
  let fname = "shaders_extracted/$1.shader" %
    (if finfo.name.len == 0: $s.pathId else: finfo.name)
  fname.writeFile(finfo.content)

let finfo =
  try:
    SerializedFileHead.deserialize(reader)
  except:
    echo "Failed on offset: ", $reader.getPosition()
    reader.close()
    quit "Error!"
    var s:SerializedFileHead
    s

echo "==========="
echo "Target file"
echo "==========="
echo "File version: ", finfo.header.version
echo "Unity3d version: ", finfo.metadata.thetype.signature
echo "Num objects: ", finfo.metadata.objects.len
echo "Finished on offset: ", $reader.getPosition()
let databegin = finfo.header.dataOffset
echo "Data offset: ", databegin
var requested = newSeq[string]()
var requested_num = newSeq[int]()
var data = newSeq[FileInfo](finfo.metadata.objects.len)
var highest = 0'u32
var nobjects = newSeq[NumberedObject](finfo.metadata.objects.len)
for i, o in finfo.metadata.objects.pairs:
  nobjects[i] = NumberedObject(idx: i, obj: o)
let soobjects = nobjects.sorted(compare)
for i, o in soobjects.pairs:
  #echo("$1 - size:$2 - cid:$3 - tid:$4" %
  #     [$o.pathId, $o.byteSize, $o.classId, $o.typeId])
  if o.obj.typeId == 48:
    let fileinfo = get_file(reader, databegin, o.obj)
    if fileinfo.name.len > 0:
      requested.add(fileinfo.name)
      requested_num.add(o.idx)
if sourcefile.len == 0:
  quit("No source file... Quitting!")
let sinfo =
  try:
    SerializedFileHead.deserialize(sreader)
  except:
    echo "Deserialization failed!"
    let pos = sreader.getPosition()
    echo "Position: ", pos
    var n:SerializedFileHead
    n
echo "==========="
echo "Source file"
echo "==========="
echo "File version: ", sinfo.header.version
echo "Unity3d version: ", sinfo.metadata.thetype.signature
echo "Num objects: ", sinfo.metadata.objects.len
var to_replace = newSeq[tuple[a:string, b:FileInfo,s:uint32]]()
var to_replace_idx = newSeq[int]()
let sdatabegin = sinfo.header.dataOffset
var ssoobjects = newSeq[NumberedObject](sinfo.metadata.objects.len)
for i, o in sinfo.metadata.objects.pairs:
  ssoobjects[i] = NumberedObject(idx:i,obj:o)
for i, o in ssoobjects.pairs:
  if o.obj.typeId == 48:
    let fileinfo = get_file(sreader, sdatabegin, o.obj)
    if fileinfo.name.len > 0 and fileinfo.name in requested:
      echo "Found: ", fileinfo.name
      let index = requested.find(fileinfo.name)
      sreader.setPosition((sdatabegin+ o.obj.byteStart).int)
      let realsize =
        if i < ssoobjects.len-1:
          ssoobjects[i+1].obj.byteStart - o.obj.byteStart
        else:
          o.obj.byteSize
      let data = sreader.readStr(realsize.int)
      to_replace.add((data, fileinfo, o.obj.byteSize))
      to_replace_idx.add(requested_num[index])

var ofile = finfo
let writer = newFileStream(outfile, fmWrite)
let tmpdata = reader.readStr(databegin.int)
writer.write(tmpdata)
while writer.getPosition() < databegin.int:
  debugEcho "Skippy"
  writer.write('\0')
var index = 0'u32
for i, o in soobjects.pairs:
  ofile.metadata.objects[o.idx].byteStart = index
  if o.idx in to_replace_idx:
    let tpl = to_replace[to_replace_idx.find(o.idx)]
    let file = tpl[0]
    let dfile = tpl[1]
    let os = tpl[2]
    debugEcho "Replacing: ", dfile.name
    writer.write(file)
    ofile.metadata.objects[o.idx].byteSize = os
  else:
  #block:
    reader.setPosition(int(databegin + o.obj.byteStart))
    let realsize =
      if i < soobjects.len-1:
        soobjects[i+1].obj.byteStart - o.obj.byteStart
      else:
        o.obj.byteSize
    let data = reader.readStr(realsize.int)
    writer.write(data)
  index = writer.getPosition().uint32 - databegin
while not reader.atEnd():
  writer.write(reader.readChar())
let fileend = writer.getPosition()
ofile.header.fileSize = fileend.uint32
writer.setPosition(0)
ofile.serialize(writer)  
sreader.close()
reader.close()
writer.close()

layout.nim

from nesm import serializable
serializable:
  type
    ObjectInfo* = tuple
      align: array[3,byte]
      pathId: int64
      byteStart: uint32
      byteSize: uint32
      typeId: int32
      classId: int16
      scriptTypeIndex: int16
      stripped: bool

serializable:
  type
    SerializedFileHeader* = object
      set: {endian: bigEndian}
      metadataSize*: uint32
      fileSize*: uint32
      version*: uint32
      dataOffset*: uint32
      set: {endian: littleEndian}
      endianess*: uint8
      reserved*: array[3, char]

    FileInfo* = tuple
      name: string
      content: string
      unknowndata: array[12, byte]

    UnityHash* = array[16, byte]

    UnityGUID* = tuple
      most: int64
      least: int64

    RTTIBaseCD* = object
      otherPartsOfId*: array[3, uint8]
      case classId*: uint8
      of 0..high(int8).uint8:
        discard        
      else:
        scriptId*: UnityHash
      oldhash*: UnityHash

    ObjectId* = tuple
      fileIndex: int32
      idInFile: int64
    
    FileId* = object
      assetPath*: cstring
      guid*: UnityGUID
      thetype*: int32
      path*: cstring
    
    RTTIDescriptor* = object
      signature*: array[8, char]
      attributes*: int32
      case embedded*: bool
      of false:
        baseClassDescriptors*: seq[RTTIBaseCD]
      else: # This surely leads to an error, but I dont care
        discard
      
    SerializedFileMetadata* = object
      thetype*: RTTIDescriptor
      objects*: seq[ObjectInfo] # align 4
      objectIds*: seq[ObjectId] # align 4!!!# Unity 5+
      padding: array[3, byte]
      externals*: seq[FileId]

    SerializedFileHead* = object
      header*: SerializedFileHeader
      metadata*: SerializedFileMetadata