moto-flash-data/commands/mpar-pax-replace.go

package commands

import (
	"archive/tar"
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	"git.cheetah.cat/cheetah/moto-flash-data/flashpart"
	"github.com/rs/zerolog/log"
)

type MparPaxReplaceCommand struct {
	InputFilename       string `long:"mpar" short:"m" required:"true" description:"Source Partition File"`
	OutputFilename      string `long:"output" short:"o" required:"true" description:"Target Partition File"`
	FindFileName        string `long:"file" short:"f" required:"true" description:"Name of the File to Find"`
	ReplaceWithFilename string `long:"replacewith" short:"r" required:"true" description:"Source Bytes for Replace-Action"`
}

func (command *MparPaxReplaceCommand) Execute(args []string) error {
	inputFile, err := os.Open(command.InputFilename)
	if err != nil {
		return err
	}
	defer inputFile.Close()
	replaceWithFile, err := os.Open(command.ReplaceWithFilename)
	if err != nil {
		return err
	}
	defer replaceWithFile.Close()

	inputFileStat, err := inputFile.Stat()
	if err != nil {
		return err
	}

	log.Info().Int64("fileSize", inputFileStat.Size()).Msg("Reading raw file...")
	pdata := make([]byte, inputFileStat.Size())
	_, err = io.ReadFull(inputFile, pdata)
	if err != nil {
		return err
	}

	fpHeader := flashpart.ParseFlashPartHeader(pdata[:32])
	log.Info().Msg(
		fmt.Sprintf("Media-ID: 0x%02X, Partition Size: 0x%08X / %d bytes, Partition Tag: %s",
			fpHeader.MediaID,
			fpHeader.TotalSize,
			fpHeader.TotalSize,
			fpHeader.VersionText,
		),
	)
	log.Info().Msg(
		fmt.Sprintf("Original-Full-Header-Hex: %s", hex.EncodeToString(pdata[:32])),
	)

	if len(pdata[32:]) != len(pdata)-32 {
		return errors.New("something doesnt add up")
	}
	if len(pdata[32:]) != int(fpHeader.TotalSize)-32 {
		log.Error().Uint32("header ts", fpHeader.TotalSize).Uint32("len pdata", uint32(len(pdata))).Msg("size mismatch")
		return nil
	}

	breader := bytes.NewReader(pdata[32:])
	tarReader := tar.NewReader(breader)

	temporaryPaxFile, _ := os.CreateTemp("", "pax-replace.tar")
	defer temporaryPaxFile.Close()
	tarWriter := tar.NewWriter(temporaryPaxFile)
	defer tarWriter.Close()

	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Error().Err(err)
			return err
		}
		if header.Typeflag == tar.TypeReg && !strings.Contains(header.Name, "._") && header.Name == command.FindFileName {
			log.Info().Str("fileName", header.Name).Int64("fileSize", header.Size).Msg("found file to modify")
			replaceWithFileStat, err := replaceWithFile.Stat()
			if err != nil {
				return err
			}
			header.Size = replaceWithFileStat.Size()
			header.ModTime = replaceWithFileStat.ModTime()
			tarWriter.WriteHeader(header)
			content, err := io.ReadAll(replaceWithFile)
			if err != nil {
				log.Error().Err(err)
			}
			tarWriter.Write(content)
		} else {
			tarWriter.WriteHeader(header)
			content, err := io.ReadAll(tarReader)
			if err != nil {
				log.Error().Err(err)
			}
			tarWriter.Write(content)
		}
	}
	//log.Error().Msg("no file replaced")

	// create new mpar with same header base
	_, err = temporaryPaxFile.Seek(0, 0)
	if err != nil {
		return err
	}
	temporaryPaxFileStat, err := temporaryPaxFile.Stat()
	if err != nil {
		return err
	}
	rawData := make([]byte, temporaryPaxFileStat.Size())
	_, err = io.ReadFull(temporaryPaxFile, rawData)
	if err != nil {
		return err
	}
	outputFile, err := os.Create(command.OutputFilename)
	if err != nil {
		return err
	}
	defer outputFile.Close()

	fpHeader.AdjustRawSize(uint32(len(rawData)))

	log.Info().Msg(
		fmt.Sprintf("New-Full-Header-Hex:      %s", hex.EncodeToString(fpHeader.GetBytes())),
	)

	fpHeaderNew := flashpart.ParseFlashPartHeader(fpHeader.GetBytes())
	log.Info().Msg(
		fmt.Sprintf("Media-ID: 0x%02X, Partition Size: 0x%08X / %d bytes, Partition Tag: %s",
			fpHeaderNew.MediaID,
			fpHeaderNew.TotalSize,
			fpHeaderNew.TotalSize,
			fpHeaderNew.VersionText,
		),
	)

	_, err = outputFile.Write(fpHeaderNew.GetBytes())
	if err != nil {
		return err
	}
	_, err = outputFile.Write(rawData)
	if err != nil {
		return err
	}

	return nil
}