123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362 |
- package main
- import (
- "flag"
- "fmt"
- "github.com/diskfs/go-diskfs/disk"
- "github.com/diskfs/go-diskfs/filesystem"
- "github.com/diskfs/go-diskfs/filesystem/fat32"
- "io"
- "io/fs"
- "os"
- "path/filepath"
- "strconv"
- "strings"
- ack "git.swzry.com/zry/go-cli-ack"
- diskfs "github.com/diskfs/go-diskfs"
- "github.com/diskfs/go-diskfs/partition/mbr"
- humanize "github.com/dustin/go-humanize"
- "github.com/liushuochen/gotable"
- yaml "gopkg.in/yaml.v3"
- )
- const (
- BLOCK_SIZE = 512
- UBOOT_OFFSET = 8192
- FIRST_PART_OFFSET = 1024 * 1024
- LAST_RESERVED = 4 * 1024 * 1024
- )
- type ConfigDef struct {
- UBootImg string `yaml:"u_boot_img"`
- BootDir string `yaml:"boot_dir"`
- BootFsSize int64 `yaml:"boot_fs_size"`
- RootImg string `yaml:"root_img"`
- UpperImg string `yaml:"upper_img"`
- SwapSize int64 `yaml:"swap_size"`
- }
- var OutputImg string
- var PrintHelp bool
- var ConfigFile string
- var Cfg ConfigDef
- var QuietMode bool
- func main() {
- flag.BoolVar(&PrintHelp, "h", false, "print this help")
- flag.StringVar(&OutputImg, "o", "sdcard.img", "output image file")
- flag.StringVar(&ConfigFile, "i", "config.yaml", "input config file")
- flag.BoolVar(&QuietMode, "q", false, "quiet mode, which does not ask the user whether to continue")
- flag.Parse()
- if PrintHelp {
- flag.PrintDefaults()
- return
- }
- fcfg, err := os.ReadFile(ConfigFile)
- if err != nil {
- panic(fmt.Sprint("failed read config file: ", err))
- }
- err = yaml.Unmarshal(fcfg, &Cfg)
- if err != nil {
- panic(fmt.Sprint("failed parse yaml config: ", err))
- }
- if Cfg.SwapSize < 4 {
- panic("swap size should not less than 4MB.")
- }
- if Cfg.BootFsSize < 4 {
- panic("swap size should not less than 4MB.")
- }
- ubootLen, err := CheckSize("u-boot", Cfg.UBootImg)
- if err != nil {
- panic(err)
- }
- ubootMaxSize := int64(FIRST_PART_OFFSET - UBOOT_OFFSET)
- if ubootLen > ubootMaxSize {
- panic(fmt.Sprintf(
- "The maximum size of u-boot is %d bytes, but yourimage has %d bytes, which exceeds %d bytes.",
- ubootMaxSize, ubootLen, ubootMaxSize-ubootLen,
- ))
- }
- CheckDir(Cfg.BootDir)
- rootLen, err := CheckSize("root filesystem", Cfg.RootImg)
- if err != nil {
- panic(err)
- }
- upperLen, err := CheckSize("upper filesystem", Cfg.UpperImg)
- if err != nil {
- panic(err)
- }
- bootActualSize := Cfg.BootFsSize * 1024 * 1024
- swapActualSize := Cfg.SwapSize * 1024 * 1024
- p1start := AlignToBlockSize(FIRST_PART_OFFSET)
- p1size := AlignToBlockSize(bootActualSize)
- p2start := p1start + p1size
- p2size := AlignToBlockSize(rootLen)
- p3start := p2start + p2size
- p3size := AlignToBlockSize(swapActualSize)
- p4start := p3start + p3size
- p4size := AlignToBlockSize(upperLen)
- diskSize := int64(p4start+p4size)*BLOCK_SIZE + LAST_RESERVED
- fmt.Println("==== Partition Table Brief ====")
- PrintPartitionTableBrief([]PTBrief{
- {
- StartSector: int(p1start),
- SizeSector: int(p1size),
- ShowName: "boot",
- },
- {
- StartSector: int(p2start),
- SizeSector: int(p2size),
- ShowName: "root",
- },
- {
- StartSector: int(p3start),
- SizeSector: int(p3size),
- ShowName: "swap",
- },
- {
- StartSector: int(p4start),
- SizeSector: int(p4size),
- ShowName: "data",
- },
- })
- fmt.Println("Total Size: ", humanize.IBytes(uint64(diskSize)))
- fmt.Println("================================")
- if !QuietMode {
- if !ack.Acknowledge("Continue?", true) {
- fmt.Println("User canceled.")
- return
- }
- }
- if !CheckExists() {
- return
- }
- idisk, err := diskfs.Create(OutputImg, diskSize, diskfs.Raw, diskfs.SectorSizeDefault)
- table := &mbr.Table{
- Partitions: []*mbr.Partition{
- {
- Bootable: false,
- Type: mbr.Fat32LBA,
- Start: p1start,
- Size: p1size,
- },
- {
- Bootable: false,
- Type: mbr.Linux,
- Start: p2start,
- Size: p2size,
- },
- {
- Bootable: false,
- Type: mbr.LinuxSwap,
- Start: p3start,
- Size: p3size,
- },
- {
- Bootable: false,
- Type: mbr.Linux,
- Start: p4start,
- Size: p4size,
- },
- },
- LogicalSectorSize: BLOCK_SIZE,
- PhysicalSectorSize: BLOCK_SIZE,
- }
- err = idisk.Partition(table)
- if err != nil {
- panic(fmt.Sprint("failed create disk partition: ", err))
- }
- fmt.Println("processing Partition 1: boot...")
- CreateFat32Fs(idisk, 1, Cfg.BootDir)
- fmt.Println("processing Partition 2: root...")
- CopyImgToPartition(idisk, 2, Cfg.RootImg)
- fmt.Println("processing Partition 3: swap...")
- FillPartition(idisk, 3, int(swapActualSize))
- fmt.Println("processing Partition 4: data...")
- CopyImgToPartition(idisk, 4, Cfg.UpperImg)
- fmt.Println("writing bootloader...")
- WriteUboot()
- fmt.Println("==== Done ====")
- }
- func CheckExists() bool {
- stat, err := os.Stat(OutputImg)
- if err != nil {
- return true
- }
- if !QuietMode {
- if !ack.Acknowledge(fmt.Sprintf("'%s' already exists, override?", OutputImg), true) {
- fmt.Println("User canceled.")
- return false
- }
- }
- if stat.IsDir() {
- panic(fmt.Sprintf("'%s' is a existed dir", OutputImg))
- }
- err = os.Remove(OutputImg)
- if err != nil {
- panic(fmt.Sprintf("'%s' is a existed file, but can not remove: %v", OutputImg, err))
- }
- return true
- }
- func AlignToBlockSize(size int64) uint32 {
- remain := size % BLOCK_SIZE
- m := size - remain
- mr := m / BLOCK_SIZE
- if remain > 0 {
- mr = mr + 1
- }
- return uint32(mr)
- }
- func CheckSize(dname, fname string) (int64, error) {
- fstat, err := os.Stat(fname)
- if err != nil {
- return 0, fmt.Errorf("can not stat %s image '%s': %v", dname, fname, err)
- }
- l := fstat.Size()
- if l <= 0 {
- return 0, fmt.Errorf("%s image is empty", dname)
- }
- return l, nil
- }
- func CheckDir(dname string) {
- fstat, err := os.Stat(dname)
- if err != nil {
- panic(fmt.Sprintf("can not stat dir '%s': %v", dname, err))
- }
- if !fstat.IsDir() {
- panic(fmt.Sprintf("'%s' is not a dir", dname))
- }
- }
- type PTBrief struct {
- ShowName string
- StartSector int
- SizeSector int
- }
- func PrintPartitionTableBrief(brief []PTBrief) {
- tb, err := gotable.Create("#", "Used For", "Offset(Sectors)", "Offset(Bytes)", "Size(Sectors)", "Size(Bytes)")
- if err != nil {
- panic(err)
- }
- for i, v := range brief {
- err = tb.AddRow(map[string]string{
- "#": strconv.Itoa(i + 1),
- "Used For": v.ShowName,
- "Offset(Sectors)": strconv.Itoa(v.StartSector),
- "Offset(Bytes)": humanize.IBytes(uint64(v.StartSector) * BLOCK_SIZE),
- "Size(Sectors)": strconv.Itoa(v.SizeSector),
- "Size(Bytes)": humanize.IBytes(uint64(v.SizeSector) * BLOCK_SIZE),
- })
- if err != nil {
- panic(err)
- }
- }
- fmt.Println(tb.String())
- }
- func CreateFat32Fs(d *disk.Disk, ptn int, from string) {
- spec := disk.FilesystemSpec{Partition: ptn, FSType: filesystem.TypeFat32, VolumeLabel: "boot"}
- ifs, err := d.CreateFilesystem(spec)
- if err != nil {
- panic(fmt.Sprintf("can not create filesystem for partition %d: %v", ptn, err))
- }
- err = filepath.Walk(from, func(path string, info fs.FileInfo, err error) error {
- if err != nil {
- fmt.Println("[Warn] error in '", path, "': ", err)
- return nil
- }
- relPath, err := filepath.Rel(from, path)
- if err != nil {
- fmt.Println("[Warn] error in '", path, "': can not get rel path: ", err)
- return nil
- }
- relPath = "/" + strings.TrimPrefix(relPath, ".")
- if info.IsDir() {
- err = ifs.Mkdir(relPath)
- if err != nil {
- fmt.Printf("[Warn] error in 'PART%d://%s': can not mkdir: %v\n", ptn, relPath, err)
- return nil
- }
- return nil
- } else {
- rw, xerr := ifs.OpenFile(relPath, os.O_CREATE|os.O_RDWR)
- if xerr != nil {
- fmt.Printf("[Warn] error in 'PART%d://%s': can not open file: %v\n", ptn, relPath, xerr)
- return nil
- }
- defer rw.Close()
- in, xerr := os.Open(path)
- if xerr != nil {
- fmt.Println("[Warn] error in '", path, "': can not open file: ", xerr)
- return nil
- }
- defer in.Close()
- _, xerr = io.Copy(rw, in)
- if xerr != nil {
- fmt.Printf("[Warn] io error in copy '%s' to 'PART%d://%s': %v\n", path, ptn, relPath, xerr)
- return nil
- }
- }
- return nil
- })
- if err != nil {
- panic(fmt.Sprint("error in copy files to partition 1: ", err))
- }
- _, ok := ifs.(*fat32.FileSystem)
- if !ok {
- panic(fmt.Errorf("create filesystem for partition %d failed: finally it is not a fat32 filesystem"))
- }
- }
- func CopyImgToPartition(d *disk.Disk, ptn int, imgfile string) {
- in, err := os.Open(imgfile)
- if err != nil {
- panic(fmt.Sprint("error in open img file '", imgfile, "': ", err))
- }
- defer in.Close()
- n, err := d.WritePartitionContents(ptn, in)
- if err != nil {
- panic(fmt.Sprintf("error in writing img '%s' to partition %d: %v '", imgfile, ptn, err))
- }
- fmt.Printf("write %s into partition %d.\n", humanize.IBytes(uint64(n)), ptn)
- }
- func FillPartition(d *disk.Disk, ptn int, size int) {
- in := NewSizableZeroReader(size)
- n, err := d.WritePartitionContents(ptn, in)
- if err != nil {
- panic(fmt.Sprintf("error in filling partition %d: %v '", ptn, err))
- }
- fmt.Printf("fill %s into partition %d.\n", humanize.IBytes(uint64(n)), ptn)
- }
- func WriteUboot() {
- fUboot, err := os.OpenFile(Cfg.UBootImg, os.O_RDONLY|os.O_EXCL, 0666)
- if err != nil {
- panic(fmt.Sprintf("error in open u-boot img: %v '", err))
- }
- defer fUboot.Close()
- fDiskImg, err := os.OpenFile(OutputImg, os.O_EXCL|os.O_RDWR, 0666)
- if err != nil {
- panic(fmt.Sprintf("error in open disk img for writting bootloader: %v '", err))
- }
- defer fUboot.Close()
- n, err := fDiskImg.Seek(UBOOT_OFFSET, 0)
- if err != nil {
- panic(fmt.Sprintf("error in seeking img for writting bootloader: %v '", err))
- }
- if n != UBOOT_OFFSET {
- panic(fmt.Sprintf("error in seeking img for writting bootloader: current pos %d != UBOOT_OFFSET(%d) '", n, UBOOT_OFFSET))
- }
- n, err = io.Copy(fDiskImg, fUboot)
- if err != nil {
- panic(fmt.Sprintf("io error in writting bootloader: %v '", err))
- }
- fmt.Printf("%d bytes written.\n", n)
- }
|