github.com/devops-filetransfer/sshego@v7.0.4+incompatible/filedb.go (about)

     1  package sshego
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"runtime"
     9  	"strings"
    10  
    11  	"github.com/glycerine/greenpack/msgp"
    12  )
    13  
    14  //go:generate greenpack
    15  
    16  var hostDbKey = "host-db"
    17  
    18  type Filedb struct {
    19  	fd       *os.File
    20  	filepath string
    21  	HostDb   *HostDb `zid:"0"`
    22  }
    23  
    24  func (b *Filedb) Close() {
    25  	if b != nil && b.fd != nil {
    26  		b.fd.Close()
    27  		b.fd = nil
    28  	}
    29  }
    30  func convertSlashes(s string) string {
    31  	return strings.Replace(s, "/", "\\", -1)
    32  }
    33  
    34  func NewFiledb(filepath string) (*Filedb, error) {
    35  
    36  	if len(filepath) == 0 {
    37  		return nil, fmt.Errorf("filepath must not be empty string")
    38  	}
    39  	if filepath[0] != '/' && filepath[0] != '.' {
    40  		if runtime.GOOS != "windows" {
    41  			// help back-compat with old prefix style argument
    42  			filepath = "./" + filepath
    43  		}
    44  	}
    45  	if runtime.GOOS == "windows" {
    46  		filepath = convertSlashes(filepath)
    47  	}
    48  
    49  	b := &Filedb{
    50  		filepath: filepath,
    51  	}
    52  	sz := int64(0)
    53  	if fileExists(filepath) {
    54  		fi, err := os.Stat(filepath)
    55  		if err != nil {
    56  			return nil, err
    57  		}
    58  		sz = fi.Size()
    59  		_ = sz
    60  	}
    61  
    62  	// maybe windows doesn't report the size?
    63  	//if sz == 0 {
    64  	//return nil, fmt.Errorf("database file present but empty! '%v'", filepath)
    65  	//}
    66  
    67  	// Open the my.db data file in your current directory.
    68  	// It will be created if it doesn't exist.
    69  	fd, err := os.OpenFile(b.filepath, os.O_RDWR|os.O_CREATE, 0600)
    70  	if err != nil {
    71  		wd, _ := os.Getwd()
    72  		// probably already open by another process.
    73  		return nil, fmt.Errorf("error opening Filedb: '%v' "+
    74  			"upon trying to open path '%s' in cwd '%s'", err, filepath, wd)
    75  	}
    76  	if err != nil {
    77  		return nil, err
    78  	}
    79  	defer fd.Close()
    80  	err = msgp.Decode(fd, b)
    81  
    82  	if err != nil {
    83  		return nil, err
    84  	}
    85  	p("FILEDB opened successfully '%s'", filepath)
    86  
    87  	return b, nil
    88  }
    89  
    90  func (b *Filedb) SaveToDisk() error {
    91  	p("Filedb.SaveToDisk is saving to b.filepath='%s'", b.filepath)
    92  
    93  	fd, err := os.OpenFile(b.filepath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
    94  	if err != nil {
    95  		return err
    96  	}
    97  	fdJson, err := os.OpenFile(b.filepath+".json", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
    98  	if err != nil {
    99  		return err
   100  	}
   101  	defer fdJson.Close()
   102  	defer fd.Close()
   103  
   104  	by, err := b.MarshalMsg(nil)
   105  	if err != nil {
   106  		return err
   107  	}
   108  	src := bytes.NewBuffer(by)
   109  
   110  	_, err = msgp.CopyToJSON(fdJson, src)
   111  	if err != nil {
   112  		return err
   113  	}
   114  	err = writeFull(fd, by)
   115  	if err != nil {
   116  		return err
   117  	}
   118  	return nil
   119  }
   120  
   121  func writeFull(w io.Writer, b []byte) error {
   122  	totw := 0
   123  	n := len(b)
   124  	for totw < n {
   125  		nw, err := w.Write(b[totw:])
   126  		if err != nil {
   127  			panic(err)
   128  			return err
   129  		}
   130  		totw += nw
   131  	}
   132  	return nil
   133  }
   134  
   135  func (b *Filedb) storeHostDb(h *HostDb) error {
   136  	b.HostDb = h
   137  	return b.SaveToDisk()
   138  }