add MBR

.travis.yml

@@ -0,0 +1,10 @@
+language: go
+
+notifications:
+  webhooks:
+    on_success: change
+    on_failure: always
+    on_start: never
+  email:
+    - tiago4orion@gmail.com
+    - tiagokatcipis@gmail.com

cmd/disk/cmds/mbr.go

@@ -0,0 +1,89 @@
+package cmds
+
+import (
+	"flag"
+	"fmt"
+
+	"github.com/barsoomia/disk/mbr"
+)
+
+var (
+	flags                       *flag.FlagSet
+	help, create, update        *bool
+	addPart, delPart, startSect *int
+	bootcode, lastSect          *string
+)
+
+func init() {
+	flags = flag.NewFlagSet("mbr", flag.ContinueOnError)
+	help = flags.Bool("help", false, "Show this help")
+	create = flags.Bool("create", false, "Create new MBR")
+	update = flags.Bool("update", false, "Update MBR")
+	addPart = flag.Int("add-part", 0, "Add partition")
+	delPart = flag.Int("del-part", 0, "Delete partition")
+	startSect = flag.Int("start-sect", 0, "start sector")
+	lastSect = flag.String("last-sect", "", "last sector (modififers +K, +M, +G works)")
+	bootcode = flags.String("bootcode", "", "Bootsector binary code")
+}
+
+func addPartition(disk string, partn int, startsect int, lastsect string) error {
+	mbrdata, err := mbr.FromFile(disk)
+	if err != nil {
+		return err
+	}
+
+	p := mbr.NewEmptyPartition()
+	mbrdata.SetPart(partn, p)
+	return nil
+}
+
+func MBR(args []string) error {
+	flags.Parse(args[1:])
+
+	if *help {
+		flags.PrintDefaults()
+		return nil
+	}
+
+	disks := flags.Args()
+	if len(disks) != 1 {
+		return fmt.Errorf("Require one device file")
+	}
+
+	if *create {
+		if *update {
+			return fmt.Errorf("-create conflicts with -update")
+		}
+
+		return mbr.Create(disks[0], *bootcode)
+	}
+
+	if *update {
+		if *addPart <= 0 || *delPart == 0 {
+			return fmt.Errorf("-update requires flag -add-part or --del-part")
+		}
+
+		if *addPart != 0 {
+			partn := *addPart
+			if *startSect == -1 {
+				return fmt.Errorf("-add-part requires -start-sect")
+			}
+
+			if *lastSect == "" {
+				return fmt.Errorf("-add-part requires -last-sect")
+			}
+			return addPartition(disks[0], partn, *startSect, *lastSect)
+		}
+
+		return fmt.Errorf("-del-part not implemented")
+	}
+
+	for _, disk := range disks {
+		err := mbr.Info(disk)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

cmd/disk/main.go

@@ -0,0 +1,43 @@
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/barsoomia/disk/cmd/disk/cmds"
+)
+
+var (
+	perr = func(format string, args ...interface{}) (int, error) {
+		return fmt.Fprintf(os.Stderr, format, args...)
+	}
+)
+
+func usage() {
+	perr("fdisk <subcommand> [options] device/file\n")
+	perr("Subcommands:\n")
+	perr("\tmbr\n")
+	perr("\n")
+	perr("Use: fdisk <subcommand> -h for more info\n")
+}
+
+func main() {
+	var err error
+	if len(os.Args) <= 1 {
+		usage()
+		os.Exit(1)
+	}
+
+	switch os.Args[1] {
+	case "mbr":
+		err = cmds.MBR(os.Args[1:])
+	default:
+		perr("Invalid subcommand: %s\n", os.Args[1])
+		os.Exit(1)
+	}
+
+	if err != nil {
+		perr("error: %s\n", err)
+		os.Exit(1)
+	}
+}

mbr/mbr.go

@@ -0,0 +1,185 @@
+package mbr
+
+import (
+	"fmt"
+	"io"
+	"math"
+	"os"
+)
+
+type (
+	mbr [MBRSize]byte
+
+	// cylinder-head-sector
+	chs struct {
+		head, sector uint8
+		cylinder     uint16
+	}
+)
+
+const (
+	MBRSize = 0x200
+	// Classical MBR structure
+	CPart1 = 0x1be // 16 bytes each
+	CPart2 = 0x1ce
+	CPart3 = 0x1de
+	CPart4 = 0x1ee
+
+	PEntrySZ = 16
+
+	Magic1Off = 0x1fe
+	Magic2Off = 0x1ff
+
+	// MBR magic numbers
+	Magic1 = 0x55
+	Magic2 = 0xaa
+
+	// bootstrap code
+	BCOffEnd = 0x01bd
+	BCSize   = BCOffEnd
+)
+
+func NewCHS(cylinder uint16, head uint8, sector uint8) chs {
+	return chs{
+		cylinder: cylinder,
+		head:     head,
+		sector:   sector,
+	}
+}
+
+func Info(fname string) error {
+	file, err := os.Open(fname)
+	if err != nil {
+		return err
+	}
+
+	var mbr [MBRSize]byte
+	_, err = file.Read(mbr[:])
+	if err != nil && err != io.EOF {
+		return err
+	}
+
+	if mbr[Magic1Off] != Magic1 &&
+		mbr[Magic2Off] != Magic2 {
+		return fmt.Errorf("no MBR found\n")
+	}
+
+	var part *partition
+	var empty bool
+
+	if part, err, empty = NewPartition(mbr[CPart1 : CPart1+PEntrySZ]); err != nil {
+		return err
+	} else if !empty {
+		fmt.Printf("%s\n", part)
+	}
+
+	if part, err, empty = NewPartition(mbr[CPart2 : CPart2+PEntrySZ]); err != nil {
+		return err
+	} else if !empty {
+		fmt.Printf("%s\n", part)
+	}
+
+	if part, err, empty = NewPartition(mbr[CPart3 : CPart3+PEntrySZ]); err != nil {
+		return err
+	} else if !empty {
+		fmt.Printf("%s\n", part)
+	}
+
+	if part, err, empty = NewPartition(mbr[CPart4 : CPart4+PEntrySZ]); err != nil {
+		return err
+	} else if !empty {
+		fmt.Printf("%s\n", part)
+	}
+
+	return nil
+}
+
+func NewMBR() mbr {
+	mbr := mbr{}
+	mbr[Magic1Off] = Magic1
+	mbr[Magic2Off] = Magic2
+	return mbr
+}
+
+func FromFile(disk string) (*mbr, error) {
+	mbr := mbr{}
+	devfile, err := os.OpenFile(disk, os.O_RDWR, 0)
+	if err != nil {
+		return nil, err
+	}
+
+	n, err := devfile.Read(mbr[:])
+	if err != nil && err != io.EOF {
+		return nil, err
+	}
+
+	if n < MBRSize {
+		return nil, fmt.Errorf("MBR requires at least 512 bytes")
+	}
+
+	return &mbr, nil
+}
+
+func (m mbr) SetBootcode(bcode []byte) error {
+	if len(bcode) > BCSize {
+		return fmt.Errorf("bootcode must have less than %d bytes", BCSize)
+	}
+
+	if copied := copy(m[0:BCOffEnd], bcode[:]); copied != len(bcode) {
+		return fmt.Errorf("Failed to copy bootcode to mbr")
+	}
+
+	return nil
+}
+
+func (m mbr) SetPart(index int, part *partition) {
+	copy(m[CPart1+index*16:16], part.Bytes())
+}
+
+func Create(devfname, bootfname string) error {
+	mbr := NewMBR()
+	devfile, err := os.OpenFile(devfname, os.O_RDWR, 0)
+	if err != nil {
+		return err
+	}
+
+	if bootfname != "" {
+		bfile, err := os.Open(bootfname)
+		if err != nil {
+			return err
+		}
+
+		var bootcode [BCSize + 1]byte
+		n, err := bfile.Read(bootcode[:])
+		if err == io.EOF || n > BCSize {
+			return fmt.Errorf("bootcode must have less than %d bytes. Got %d", BCSize, n)
+		}
+
+		err = mbr.SetBootcode(bootcode[0:n])
+		if err != nil {
+			return err
+		}
+	}
+
+	_, err = devfile.Write(mbr[:])
+	return err
+}
+
+func CHS2LBA(c uint16, h uint8, s uint8) uint32 {
+	return (uint32(c)*16+uint32(h))*63 + uint32((s - 1))
+}
+
+func LBA2C(lba uint32) uint16 {
+	return uint16(math.Mod(float64(lba), 16*63))
+}
+
+func LBA2H(lba uint32) uint8 {
+	lbaf := float64(lba)
+	spt := float64(63)
+	hpt := float64(16)
+	return uint8(math.Mod(math.Mod(lbaf, spt), hpt))
+}
+
+func LBA2S(lba uint32) uint8 {
+	return uint8(math.Mod(float64(lba), float64(63))) + 1
+}

mbr/mbr_test.go

@@ -0,0 +1,50 @@
+package mbr
+
+import "testing"
+
+type testTbl struct {
+	chs chs
+	lba uint32
+}
+
+func TestCHS2LBA(t *testing.T) {
+	var tests []testTbl
+
+	for i := 1; i < 64; i++ {
+		tests = append(tests, struct {
+			chs chs
+			lba uint32
+		}{
+			chs: NewCHS(0, 0, uint8(i)),
+			lba: uint32(i - 1),
+		})
+	}
+
+	tests = append(tests, testTbl{
+		chs: NewCHS(0, 1, 1),
+		lba: 63,
+	})
+
+	tests = append(tests, testTbl{
+		chs: NewCHS(0, 15, 1),
+		lba: 945,
+	})
+
+	tests = append(tests, testTbl{
+		chs: NewCHS(15, 15, 63),
+		lba: 16127,
+	})
+
+	tests = append(tests, testTbl{
+		chs: NewCHS(16319, 15, 63),
+		lba: 16450559,
+	})
+
+	for _, test := range tests {
+		if val := CHS2LBA(test.chs.cylinder, test.chs.head, test.chs.sector); val != test.lba {
+			t.Errorf("Expected %d but got %d", test.lba, val)
+			return
+		}
+	}
+
+}