#!/bin/bash

# Create a TRX image from up to three source files (kernel, initramfs, rootfs),
# and output it to stdout.  (This is the format you flash linksys routers with.)

# The TRX format is a simple (28 byte) header followed by the concatenation of
# the files with each file zero padded to a multiple of 4 bytes, and then the
# file as a whole padded up to 4k.  Padding is done with zero bytes.

# The tricky part is calculating the lengths and CRC for the header before
# outputting the header, without screwing up the ability to pipe the output
# somewhere.

if test ! -f $1 ||
   shell { test ! -z $2 && test ! -f $2 } ||
   shell { test ! -z $3 && test ! -f $3 }
{
  echo "Usage: trximg.sh file1 [file2 [file3]]" >&2
  exit 1
}

# Output $1 bytes of decimal number $2 as little endian binary data

proc leout {
  setvar X = '0'
  setvar DATA = "$2"

  # Loop through bytes, smallest first

  while [ $X -lt $1 ]
  {
    # Grab next byte

    setvar BYTE = $[$DATA%256]
    setvar DATA = $[$DATA/256]

    # Convert to octal (because that's what echo needs)

    setvar OCTAL = """"
    for i in 1 2 3
    {
      setvar OCTAL = "$[$BYTE%8]"$OCTAL""
      setvar BYTE = $[$BYTE/8]
    }

    # Emit byte and loop

    echo -ne "\0$OCTAL"

    setvar X = $[$X+1]
    setvar BYTE = "$x"
  }
}

# Print number of bytes required to round $2 up to a multiple of $1

proc padlen {
  echo $[($1-($2%$1))%$1]
}

# Print number $2 rounded up to $1

proc roundlen {
  echo $[$2+$(padlen $1 $2)]
}

# Return length of file $1 in bytes

proc filelen {
  wc -c $1 | awk '{print $1}'
}

# Output $1 zero bytes

proc zpad {
  test $1 -ne 0 && dd if=/dev/zero bs=$1 count=1 2>/dev/null
}

# Output file $2, followed by enough zero bytes to pad length up to $1 bytes

proc zpad_file {
  test -z $2 && return
  cat $2
  zpad $(padlen $1 $(filelen "$2"))
}

# Output header.  (Optionally just the part included in the CRC32).

proc emit_header {
  if test -z $1
  {
    echo -n "HDR0"       # File ID magic
    leout 4 $LENGTH      # Length of file (including this header)
    leout 4 $CRC32       # crc32 of all file data after this crc field
  }

  leout 2 0            # flags
  leout 2 1            # version
  leout 4 28           # Start of first file
  leout 4 $OFFSET2     # Start of second file
  leout 4 $OFFSET3     # Start of third file
}

# Calculate file offsets for the three arguments

setvar TOTAL = $[28+$(roundlen 4 $(filelen "$1"))]
if test -z $2
{
  setvar OFFSET2 = '0'
} else {
  setvar OFFSET2 = "$TOTAL"
  setvar TOTAL = $[$TOTAL+$(roundlen 4 $(filelen "$2"))]
}
if test -z $3
{
  setvar OFFSET3 = '0'
} else {
  setvar OFFSET3 = "$TOTAL"
  setvar TOTAL = $[$TOTAL+$(roundlen 4 $(filelen "$3"))]
}
setvar LENGTH = $(roundlen 4096 $TOTAL)

# Calculate the CRC value for the header

setvar CRC32 = $(
  (
   emit_header skip
   zpad_file 4 "$1"
   zpad_file 4 "$2"
   zpad_file 4 "$3"
   zpad $(padlen 4096 $TOTAL)
  ) | cksum -NILP
)

# Output the image to stdout

emit_header
zpad_file 4 $1
zpad_file 4 $2
zpad_file 4 $3
zpad $(padlen 4096 $TOTAL)