echo 'Western Digital ROYL write module from file using vendor specific commands.'
echo 'WARNING: THIS IS DANGEROUS AND COULD KILL THE DRIVE!!!'
echo 'USE AT YOUR OWN RISK!!!'
echo 'It is IMPORTANT to use 0x to select the module number in HEX!'
echo 'The difference between decimal and hex COULD KILL THE DRIVE!!!'
echo "This modifies data in the service area of the drive!"
echo "If you don't know what you are doing, then don't do it!"

seti $printhelp = $printhelp
seti $help = $help
if $help = 1
  echo 'Western Digital ROYL write module from file using vendor specific commands.'
  echo 'Note that some USB drives do not support these vendor specific commands.'
  echo 'WARNING! THIS SCRIPT COULD BE DANGEROUS IF YOU DO NOT KNOW WHAT YOU ARE DOING!'
  echo 'It will create a folder using the model and serial as the folder name,'
  echo ' and create a backup dump file with the date and time as part of the name.'
  echo 'It will create a backup dump every time a read or write is attempted.'
  echo 'The checksum will be recalculated if it not correct.'
  echo 'This requires the number variable "mod" to be set in DECIMAL.'
  echo 'To enter a HEX number preceed it with "0x" (IMPORTANT TO SELECT CORRECT MODULE).'
  echo '"mod" is the SA module to be written.'
  echo 'This requires the text variable "file" to be set.'
  echo '"file" is the file which is to be written to the module.'

  if $printhelp = 1
    exit 0
  endif
  echo 'Hit ENTER to continue...'
  userinput $choicestring
  previousscript
endif

include good_subroutines

# set a blank string variable to compare with user input
# make sure not to set this anyplace else
sets $null = ""

# set this to itself for use with the menu system
# make sure not to set this anyplace else
# if it was not previously set it will = 0
seti $using_menu = $using_menu

# set the module number
seti $ask = 0
while 1 = 1
  variablecheck $mod
  if $error_level < 16
    seti $ask = 1
  elseif $using_menu = 1
    seti $ask = 1
  endif
  if $ask = 1
    echo "Enter module number in decimal (0x for hex):"
    userinput $choicestring
    if $choicestring != $null
      seti $mod = $choicestring 0
      seti $module_number = $mod
      break
    else
      echo "Choice cannot be blank!"
    endif
  else
    break
  endif
done

# set the filename to be read
while 1 = 1
  variablecheck $file
  if $error_level < 16
    seti $ask = 1
  elseif $using_menu = 1
    seti $ask = 1
  endif
  if $ask = 1
    echo "Enter file name:"
    userinput $choicestring
    if $choicestring != $null
      sets $file = $choicestring
      sets $patchname = $file
      break
    else
      echo "Choice cannot be blank!"
    endif
  else
    break
  endif
done

# set the filenames to be written to
hex
sets $basefilename = "mod0x" $module_number "original.bin"
sets $basepatchname = "mod0x" $module_number "patched.bin"
decimal

echo "WARNING! You are about to write to the service area of the drive."
echo "You are on your own if something goes wrong."
echo 'Type "DANGEROUS" to continue:'
userinput $choicestring
sets $safety = "DANGEROUS"
if $choicestring != $safety
  echo "You chose not to continue"
  previousscript
endif

gosub write_patch


previousscript
end






subroutine read_module
  # get model and serial using identify device command
  echo "Performing identify device command"
  buffersize 512
  setreadpio
  ata28cmd 0 0 0 0 0 0xa0 0xec
  # check if command failed
  gosub status_check
  # extract model and serial and trim leading and trailing spaces
  sets $model = "null"
  sets $serial = "null"
  wordflipbuffer 0 512
  seti $count = 0
  seti $start_offset = 54
  while $count < 40
    seti $byte = buffer $start_offset
    if $byte > 0x20
      break
    endif
    seti $start_offset = $start_offset + 1
    seti $count = $count + 1
  done
  seti $count = 0
  seti $end_offset = 93
  while $count < 40
    seti $byte = buffer $end_offset
    if $byte > 0x20
      break
    endif
    seti $end_offset = $end_offset - 1
    seti $count = $count + 1
  done
  seti $end_offset = $end_offset + 1
  seti $size = $end_offset - $start_offset
  if $size > 0
    sets $model = buffer $start_offset $size
  endif
  # find out if it is a WD drive
  sets $wdcheck = buffer $start_offset 3
  sets $compare = "WDC"
  if $wdcheck != $compare
    echo "Model: " $model
    echo "Drive is not WD, exiting
    previousscript
  endif

  seti $count = 0
  seti $start_offset = 20
  while $count < 20
    seti $byte = buffer $start_offset
    if $byte > 0x20
      break
    endif
    seti $start_offset = $start_offset + 1
    seti $count = $count + 1
  done
  seti $count = 0
  seti $end_offset = 39
  while $count < 20
    seti $byte = buffer $end_offset
    if $byte > 0x20
      break
    endif
    seti $end_offset = $end_offset - 1
    seti $count = $count + 1
  done
  seti $end_offset = $end_offset + 1
  seti $size = $end_offset - $start_offset
  if $size > 0
    sets $serial = buffer $start_offset $size
  endif

  clearbuffer

  echo "Model: " $model
  echo "Serial: " $serial

  # create a directory named with model and serial
  sets $directory = $model "__" $serial
  sets $command = 'mkdir -p "' $directory '"'
  callcommand $command

  # set filename for backup dump file
  gettime
  sets $backup_filename = $directory "/" $date "__" $basefilename
  sets $backup_patchname = $directory "/" $date "__" $basepatchname


  echo "Reading the module"

  gosub enable_vsc

  gosub prepare_read_vsc

  # vsc to get a sector of the module
  setreadpio
  buffersize 0x200
  clearbuffer
  ata28cmd 0xd5 0x01 0xbf 0x4f 0xc2 0xa0 0xb0
  # check if command failed
  gosub status_check

  # find how many sectors the module contains
  seti $mod_length_sectors = buffer 0xa
  seti $tempnum = buffer 0xb
  seti $tempnum = $tempnum > 8
  seti $mod_length_sectors = $mod_length_sectors + $tempnum
  seti $remaining_sectors = $mod_length_sectors - 1

  # copy the sector to the scratchpad
  seti $scratchpad_size = $mod_length_sectors * 512
  scratchpadsize $scratchpad_size
  copybuffertoscratchpad 0 0 0x200

  # vsc to get the remaining sectors
  # this gets them one at a time
  # there is a way to get them all at once,
  # but there was a case where that did not work
  seti $count = 1
  while $count < $mod_length_sectors
    setreadpio
    buffersize 0x200
    clearbuffer
    ata28cmd 0xd5 0x01 0xbf 0x4f 0xc2 0xa0 0xb0
    # check if command failed
    gosub status_check
    # copy the sectors to the scratchpad
    seti $offset = $count * 0x200
    copybuffertoscratchpad 0 $offset 0x200
    seti $count = $count + 1
  done

  gosub disable_vsc

  # write the sectors to the backup file
  writescratchpad $backup_filename 0 0 $scratchpad_size

  # process the module
  echo "Header:"
  printscratchpad 0 4
  seti $mod_id = scratchpad 8 w
  hex
  echo "Module ID = 0x" $mod_id
  echo "Size in sectors = 0x" $mod_length_sectors
  seti $checksum = scratchpad 0xc dw
  echo "32 bit checksum = 0x" $checksum
  echo "Mod version:"
  printscratchpad 0x10 8
  decimal
  sets $header = scratchpad 0 4
  sets $header_check = "ROYL"
  if $header != $header_check
    echo "Header is not 'ROYL', exiting"
    previousscript
  endif

  # calculate checksum
  seti $calculated_checksum = 0
  gosub calculate_checksum_scratchpad
  hex
  if $checksum = $calculated_checksum
    echo "Calculated Checksum = 0x" $calculated_checksum " (good)"
  else
    echo "Calculated Checksum = 0x" $calculated_checksum " (BAD)"
    echo "Checksum of read module is not valid, exiting"
    previousscript
  endif
  decimal

endsubroutine





subroutine make_patch
  # this patches the checksum
  # calculate new checksum and insert it
  seti $calculated_checksum = 0
  gosub calculate_checksum_scratchpad
  hex
  echo "Calculated Checksum = 0x" $calculated_checksum
  decimal
  seti $b1 = $calculated_checksum & 0xff
  seti $calculated_checksum = $calculated_checksum > 8
  seti $b2 = $calculated_checksum & 0xff
  seti $calculated_checksum = $calculated_checksum > 8
  seti $b3 = $calculated_checksum & 0xff
  seti $calculated_checksum = $calculated_checksum > 8
  seti $b4 = $calculated_checksum & 0xff
  setscratchpad 0x0c
    $b1 $b2 $b3 $b4
  endscratchpad
endsubroutine






subroutine status_check
  seti $command_failed = 0
  gosub check_command_status
  if $command_failed = 1
    echo "Command failed!"
    gosub show_sense_data
    gosub show_ata_return_status
    previousscript
  endif
endsubroutine





subroutine enable_vsc
  # enable vendor specific commands
  echo "enable vsc"
  buffersize 0
  setreadpio
  ata28cmd 0x45 0x0b 0x00 0x44 0x57 0xa0 0x80
  # check if command failed
  gosub status_check
endsubroutine






subroutine disable_vsc
  # disable vendor specific commands
  echo "disable vsc"
  buffersize 0
  setreadpio
  ata28cmd 0x44 0x0b 0x00 0x44 0x57 0xa0 0x80
  # check if command failed
  gosub status_check
endsubroutine






subroutine write_patch
  gosub read_module
  getfilesize $patchname
  seti $patch_size = $error_level
  if $patch_size < 0
    echo "ERROR! File " $patchname " not found!"
    echo "The patch has not been written to the drive!"
    returnsub
  endif
  seti $divide_check = $patch_size % 512
  if $divide_check != 0
    echo "ERROR! Filesize of " $patchname " is not evenly dividable by sector size!"
    echo "The patch has not been written to the drive!"
    returnsub
  endif

  echo "Reading the patch file"
  seti $scratchpad_size = $patch_size
  scratchpadsize $scratchpad_size
  clearscratchpad
  readscratchpad $patchname 0 0 $scratchpad_size

  # find how many sectors the module contains
  seti $mod_length_sectors = scratchpad 0xa
  seti $tempnum = scratchpad 0xb
  seti $tempnum = $tempnum > 8
  seti $mod_length_sectors = $mod_length_sectors + $tempnum

  # process the module
  echo "Header:"
  printscratchpad 0 4
  seti $mod_id = scratchpad 8 w
  hex
  echo "Module ID = 0x" $mod_id
  echo "Size in sectors = 0x" $mod_length_sectors
  seti $checksum = scratchpad 0xc dw
  echo "32 bit checksum = 0x" $checksum
  echo "Mod version:"
  printscratchpad 0x10 8
  decimal
  sets $header = scratchpad 0 4
  sets $header_check = "ROYL"
  if $header != $header_check
    echo "Header is not 'ROYL', exiting"
    previousscript
  endif

  seti $size_check = $patch_size / 512
  if $size_check != $mod_length_sectors
    echo "ERROR! The size of " $patchname " does not match the value read from the module!"
    echo "The patch has not been written to the drive!"
    returnsub
  endif

  if $module_number != $mod_id
    echo "ERROR! The module number does not match!"
    echo "The patch has not been written to the drive!"
    returnsub
  endif

  # verify checksum
  seti $checksum = scratchpad 0xc dw
  seti $calculated_checksum = 0
  gosub calculate_checksum_scratchpad
  hex
  if $checksum != $calculated_checksum
    echo "Checksum in patch file = 0x" $checksum
    echo "Calculated Checksum = 0x" $calculated_checksum
    echo "Checksum of patch file is not valid."
    echo "Fixing the Checksum."
    gosub make_patch
    seti $checksum = scratchpad 0xc dw
    hex
    echo "New checksum = 0x" $checksum
  endif
  decimal

  echo ""
  echo "Writing the patched module to the drive"

  gosub enable_vsc

  gosub prepare_write_vsc

  # vsc to write the sectors to the disk
  # this writes them one at a time
  # there is a way to get them all at once,
  # but there was a case where that did not work
  seti $count = 0
  echo $count
  while $count < $mod_length_sectors
    setwritepio
    buffersize 0x200
    # copy the sector to the buffer
    seti $offset = $count * 0x200
    copyscratchpadtobuffer $offset 0 0x200
    ata28cmd 0xd6 0x01 0xbf 0x4f 0xc2 0xa0 0xb0
    # check if command failed
    gosub status_check
    seti $count = $count + 1
    echo $count
  done

  gosub disable_vsc

  # read the module again and compare it to what we wrote
  echo "Verifying data"

  gosub enable_vsc

  gosub prepare_read_vsc

  seti $verify_check = 0
  gosub read_and_verify

  if $verify_check = 0
    echo "Data verified!"
  else
    echo "WARNING! The data failed verification!"
    echo "It appears the data write was unsuccessful!"
  endif

  # write the sectors to the backup file
  writescratchpad $backup_patchname 0 0 $scratchpad_size

  gosub disable_vsc

endsubroutine







subroutine prepare_read_vsc
  # vsc to prepare (command to set for module read)
  echo "prepare read vsc"
  buffersize 0x200
  clearbuffer
  seti $idlow = $module_number & 0xff
  seti $idhigh = $module_number > 8
  setbuffer 0
    08 00 01 00 $idlow $idhigh
  endbuffer
  setwritepio
  ata28cmd 0xd6 0x01 0xbe 0x4f 0xc2 0xa0 0xb0
  # check if command failed
  gosub status_check
endsubroutine




subroutine prepare_write_vsc
  # vsc to prepare (command to set for module write)
  echo "prepare write vsc"
  buffersize 0x200
  clearbuffer
  seti $idlow = $module_number & 0xff
  seti $idhigh = $module_number > 8
  setbuffer 0
    08 00 02 00 $idlow $idhigh
  endbuffer
  setwritepio
  ata28cmd 0xd6 0x01 0xbe 0x4f 0xc2 0xa0 0xb0
  # check if command failed
  gosub status_check
endsubroutine




subroutine read_and_verify
  # vsc to get read the sectors
  # this gets them and verifies one at a time
  seti $count = 0
  while $count < $mod_length_sectors
    buffersize 0x200
    clearbuffer
    setreadpio
    ata28cmd 0xd5 0x01 0xbf 0x4f 0xc2 0xa0 0xb0
    # check if command failed
    gosub status_check
    # compare the data
    seti $offset = $count * 0x200
    seti $byte_count = 0
    while $byte_count < 256
      seti $b = buffer $byte_count
      seti $s = scratchpad $offset
      if $b != $s
        seti $verify_check = 1
        break
      endif
      seti $offset = $offset + 1
      seti $byte_count = $byte_count + 1
    done
    if $verify_check != 0
      break
    endif
    seti $count = $count + 1
  done
endsubroutine




subroutine calculate_checksum_scratchpad
seti $calculated_checksum = 0
  seti $count = 0
  while $count < $scratchpad_size
    seti $dword = scratchpad $count dw
    seti $calculated_checksum = $calculated_checksum + $dword
    seti $count = $count + 4
    # skip the actual checksum bytes
    if $count = 0xc
      seti $count = $count + 4
    endif
  done
  seti $calculated_checksum = $calculated_checksum & 0xffffffff
  seti $calculated_checksum = 0xffffffff - $calculated_checksum
  seti $calculated_checksum = $calculated_checksum + 1
endsubroutine