arozos/mod/filesystem/abstractions/smbfs/smbfs.go

package smbfs

import (
	"fmt"
	"io"
	"io/fs"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/hirochachacha/go-smb2"
	"imuslab.com/arozos/mod/filesystem/fsdef"
)

/*
	Server Message Block.go

	This is a file abstraction that mount SMB folders onto ArozOS as virtual drive

*/

type ServerMessageBlockFileSystemAbstraction struct {
	UUID      string
	Hierarchy string
	ipaddr    string
	user      string
	conn      *net.Conn
	session   *smb2.Session
	share     *smb2.Share
}

func NewServerMessageBlockFileSystemAbstraction(uuid string, hierarchy string, ipaddr string, rootShare string, username string, password string) (ServerMessageBlockFileSystemAbstraction, error) {
	conn, err := net.Dial("tcp", ipaddr)
	if err != nil {
		log.Println("[SMB-FS] Unable to connect to remote: ", err.Error())
		return ServerMessageBlockFileSystemAbstraction{}, err
	}

	d := &smb2.Dialer{
		Initiator: &smb2.NTLMInitiator{
			User:     username,
			Password: password,
		},
	}

	s, err := d.Dial(conn)
	if err != nil {
		log.Println("[SMB-FS] Unable to connect to remote: ", err.Error())
		return ServerMessageBlockFileSystemAbstraction{}, err
	}

	//Mound remote storage
	fs, err := s.Mount(rootShare)
	if err != nil {
		log.Println("[SMB-FS] Unable to connect to remote: ", err.Error())
		return ServerMessageBlockFileSystemAbstraction{}, err
	}

	log.Println("[SMB-FS] Connected to remote: " + ipaddr)
	return ServerMessageBlockFileSystemAbstraction{
		UUID:      uuid,
		Hierarchy: hierarchy,
		ipaddr:    ipaddr,
		user:      username,
		conn:      &conn,
		session:   s,
		share:     fs,
	}, nil
}

func (a ServerMessageBlockFileSystemAbstraction) Chmod(filename string, mode os.FileMode) error {
	filename = filterFilepath(filename)
	filename = toWinPath(filename)
	return a.share.Chmod(filename, mode)
}
func (a ServerMessageBlockFileSystemAbstraction) Chown(filename string, uid int, gid int) error {
	return fsdef.ErrOperationNotSupported
}
func (a ServerMessageBlockFileSystemAbstraction) Chtimes(filename string, atime time.Time, mtime time.Time) error {
	filename = filterFilepath(filename)
	filename = toWinPath(filename)
	return a.share.Chtimes(filename, atime, mtime)
}
func (a ServerMessageBlockFileSystemAbstraction) Create(filename string) (*os.File, error) {
	return nil, fsdef.ErrOperationNotSupported
}
func (a ServerMessageBlockFileSystemAbstraction) Mkdir(filename string, mode os.FileMode) error {
	filename = filterFilepath(filename)
	filename = toWinPath(filename)
	return a.share.Mkdir(filename, mode)
}
func (a ServerMessageBlockFileSystemAbstraction) MkdirAll(filename string, mode os.FileMode) error {
	filename = filterFilepath(filename)
	filename = toWinPath(filename)
	return a.share.MkdirAll(filename, mode)
}
func (a ServerMessageBlockFileSystemAbstraction) Name() string {
	return ""
}
func (a ServerMessageBlockFileSystemAbstraction) Open(filename string) (*os.File, error) {
	return nil, fsdef.ErrOperationNotSupported
}
func (a ServerMessageBlockFileSystemAbstraction) OpenFile(filename string, flag int, perm os.FileMode) (*os.File, error) {
	return nil, fsdef.ErrOperationNotSupported
}
func (a ServerMessageBlockFileSystemAbstraction) Remove(filename string) error {
	filename = filterFilepath(filename)
	filename = toWinPath(filename)
	return a.share.Remove(filename)
}
func (a ServerMessageBlockFileSystemAbstraction) RemoveAll(filename string) error {
	filename = filterFilepath(filename)
	filename = toWinPath(filename)
	return a.share.RemoveAll(filename)
}
func (a ServerMessageBlockFileSystemAbstraction) Rename(oldname, newname string) error {
	oldname = toWinPath(filterFilepath(oldname))
	newname = toWinPath(filterFilepath(newname))
	return a.share.Rename(oldname, newname)
}
func (a ServerMessageBlockFileSystemAbstraction) Stat(filename string) (os.FileInfo, error) {
	filename = toWinPath(filterFilepath(filename))
	return a.share.Stat(filename)
}
func (a ServerMessageBlockFileSystemAbstraction) Close() error {
	a.share.Umount()
	a.session.Logoff()
	conn := *(a.conn)
	conn.Close()
	return nil
}

/*
	Abstraction Utilities
*/

func (a ServerMessageBlockFileSystemAbstraction) VirtualPathToRealPath(subpath string, username string) (string, error) {
	if strings.HasPrefix(subpath, a.UUID+":") {
		//This is full virtual path. Trim the uuid and correct the subpath
		subpath = strings.TrimPrefix(subpath, a.UUID+":")
	}
	subpath = filterFilepath(subpath)

	if a.Hierarchy == "user" {
		return toWinPath(filepath.ToSlash(filepath.Clean(filepath.Join("users", username, subpath)))), nil
	} else if a.Hierarchy == "public" {
		return toWinPath(filepath.ToSlash(filepath.Clean(subpath))), nil
	}

	return "", fsdef.ErrVpathResolveFailed
}

func (a ServerMessageBlockFileSystemAbstraction) RealPathToVirtualPath(fullpath string, username string) (string, error) {
	fullpath = strings.TrimPrefix(fullpath, "\\")
	vpath := a.UUID + ":/" + filterFilepath(fullpath)
	return vpath, nil
}

func (a ServerMessageBlockFileSystemAbstraction) FileExists(realpath string) bool {
	realpath = toWinPath(filterFilepath(realpath))
	_, err := a.share.Open(realpath)
	return err == nil
}

func (a ServerMessageBlockFileSystemAbstraction) IsDir(realpath string) bool {
	realpath = filterFilepath(realpath)
	realpath = toWinPath(realpath)
	stx, err := a.share.Stat(realpath)
	if err != nil {
		return false
	}
	return stx.IsDir()
}

func (a ServerMessageBlockFileSystemAbstraction) Glob(realpathWildcard string) ([]string, error) {
	realpathWildcard = toWinPath(filterFilepath(realpathWildcard))
	return a.share.Glob(realpathWildcard)
}

func (a ServerMessageBlockFileSystemAbstraction) GetFileSize(realpath string) int64 {
	realpath = toWinPath(filterFilepath(realpath))
	stat, err := a.share.Stat(realpath)
	if err != nil {
		return 0
	}
	return stat.Size()
}

func (a ServerMessageBlockFileSystemAbstraction) GetModTime(realpath string) (int64, error) {
	realpath = toWinPath(filterFilepath(realpath))
	stat, err := a.share.Stat(realpath)
	if err != nil {
		return 0, nil
	}
	return stat.ModTime().Unix(), nil
}

func (a ServerMessageBlockFileSystemAbstraction) WriteFile(filename string, content []byte, mode os.FileMode) error {
	filename = toWinPath(filterFilepath(filename))
	return a.share.WriteFile(filename, content, mode)
}
func (a ServerMessageBlockFileSystemAbstraction) ReadFile(filename string) ([]byte, error) {
	filename = toWinPath(filterFilepath(filename))
	return a.share.ReadFile(filename)
}
func (a ServerMessageBlockFileSystemAbstraction) WriteStream(filename string, stream io.Reader, mode os.FileMode) error {
	filename = toWinPath(filterFilepath(filename))
	f, err := a.share.OpenFile(filename, os.O_CREATE|os.O_WRONLY, mode)
	if err != nil {
		return err
	}

	p := make([]byte, 1024)
	for {
		_, err := stream.Read(p)
		if err != nil {
			if err == io.EOF {
				break
			} else {
				return err
			}
		}
		_, err = f.Write(p)
		if err != nil {
			return err
		}
	}

	return nil
}
func (a ServerMessageBlockFileSystemAbstraction) ReadStream(filename string) (io.ReadCloser, error) {
	filename = toWinPath(filterFilepath(filename))
	f, err := a.share.OpenFile(filename, os.O_RDONLY, 0755)
	if err != nil {
		return nil, err
	}
	return f, nil
}

func (a ServerMessageBlockFileSystemAbstraction) Walk(root string, walkFn filepath.WalkFunc) error {
	root = toWinPath(filterFilepath(root))
	err := fs.WalkDir(a.share.DirFS(root), ".", func(path string, d fs.DirEntry, err error) error {
		statInfo, _ := d.Info()
		walkFn(path, statInfo, err)
		return nil
	})
	return err
}

/*

	Optional Functions

*/

func (a *ServerMessageBlockFileSystemAbstraction) CapacityInfo() {
	fsinfo, err := a.share.Statfs(".")
	if err != nil {
		return
	}

	fmt.Println(fsinfo)
}

/*

	Helper Functions

*/

func toWinPath(filename string) string {
	backslashed := strings.ReplaceAll(filename, "/", "\\")
	return strings.TrimPrefix(backslashed, "\\")

}

func filterFilepath(rawpath string) string {
	rawpath = filepath.ToSlash(filepath.Clean(rawpath))
	rawpath = strings.TrimSpace(rawpath)
	if strings.HasPrefix(rawpath, "./") {
		return rawpath[1:]
	} else if rawpath == "." || rawpath == "" {
		return "/"
	}
	return rawpath
}