1package tree
2
3import (
4	"errors"
5	"fmt"
6	"io"
7	"os"
8	"os/user"
9	"path/filepath"
10	"regexp"
11	"sort"
12	"strconv"
13	"strings"
14)
15
16// Node represent some node in the tree
17// contains FileInfo, and its childs
18type Node struct {
19	os.FileInfo
20	path   string
21	depth  int
22	err    error
23	nodes  Nodes
24	vpaths map[string]bool
25}
26
27// List of nodes
28type Nodes []*Node
29
30// To use this package programmatically, you must implement this
31// interface.
32// For example: PTAL on 'cmd/tree/tree.go'
33type Fs interface {
34	Stat(path string) (os.FileInfo, error)
35	ReadDir(path string) ([]string, error)
36}
37
38// Options store the configuration for specific tree.
39// Note, that 'Fs', and 'OutFile' are required (OutFile can be os.Stdout).
40type Options struct {
41	Fs      Fs
42	OutFile io.Writer
43	// List
44	All        bool
45	DirsOnly   bool
46	FullPath   bool
47	IgnoreCase bool
48	FollowLink bool
49	DeepLevel  int
50	Pattern    string
51	IPattern   string
52	MatchDirs  bool
53	Prune      bool
54	// File
55	ByteSize bool
56	UnitSize bool
57	FileMode bool
58	ShowUid  bool
59	ShowGid  bool
60	LastMod  bool
61	Quotes   bool
62	Inodes   bool
63	Device   bool
64	// Sort
65	NoSort    bool
66	VerSort   bool
67	ModSort   bool
68	DirSort   bool
69	NameSort  bool
70	SizeSort  bool
71	CTimeSort bool
72	ReverSort bool
73	// Graphics
74	NoIndent bool
75	Colorize bool
76	// Color defaults to ANSIColor()
77	Color func(*Node, string) string
78}
79
80func (opts *Options) color(node *Node, s string) string {
81	f := opts.Color
82	if f == nil {
83		f = ANSIColor
84	}
85	return f(node, s)
86}
87
88// New get path and create new node(root).
89func New(path string) *Node {
90	return &Node{path: path, vpaths: make(map[string]bool)}
91}
92
93// Visit all files under the given node.
94func (node *Node) Visit(opts *Options) (dirs, files int) {
95	// visited paths
96	if path, err := filepath.Abs(node.path); err == nil {
97		path = filepath.Clean(path)
98		node.vpaths[path] = true
99	}
100	// stat
101	fi, err := opts.Fs.Stat(node.path)
102	if err != nil {
103		node.err = err
104		return
105	}
106	node.FileInfo = fi
107	if !fi.IsDir() {
108		return 0, 1
109	}
110	// increase dirs only if it's a dir, but not the root.
111	if node.depth != 0 {
112		dirs++
113	}
114	// DeepLevel option
115	if opts.DeepLevel > 0 && opts.DeepLevel <= node.depth {
116		return
117	}
118	// MatchDirs option
119	var dirMatch bool
120	if node.depth != 0 && opts.MatchDirs {
121		// then disable prune and pattern for immediate children
122		if opts.Pattern != "" {
123			dirMatch = node.match(opts.Pattern, opts)
124		} else if opts.IPattern != "" && node.match(opts.IPattern, opts) {
125			return
126		}
127	}
128	names, err := opts.Fs.ReadDir(node.path)
129	if err != nil {
130		node.err = err
131		return
132	}
133	node.nodes = make(Nodes, 0)
134	for _, name := range names {
135		// "all" option
136		if !opts.All && strings.HasPrefix(name, ".") {
137			continue
138		}
139		nnode := &Node{
140			path:   filepath.Join(node.path, name),
141			depth:  node.depth + 1,
142			vpaths: node.vpaths,
143		}
144		d, f := nnode.Visit(opts)
145		if nnode.IsDir() {
146			// "prune" option, hide empty directories
147			if opts.Prune && f == 0 {
148				continue
149			}
150			if opts.MatchDirs && opts.IPattern != "" && nnode.match(opts.IPattern, opts) {
151				continue
152			}
153		} else if nnode.err == nil {
154			// "dirs only" option
155			if opts.DirsOnly {
156				continue
157			}
158			// Pattern matching
159			if !dirMatch && opts.Pattern != "" && !nnode.match(opts.Pattern, opts) {
160				continue
161			}
162			// IPattern matching
163			if opts.IPattern != "" && nnode.match(opts.IPattern, opts) {
164				continue
165			}
166		}
167		node.nodes = append(node.nodes, nnode)
168		dirs, files = dirs+d, files+f
169	}
170	// Sorting
171	if !opts.NoSort {
172		node.sort(opts)
173	}
174	return
175}
176
177func (node *Node) match(pattern string, opt *Options) bool {
178	var prefix string
179	if opt.IgnoreCase {
180		prefix = "(?i)"
181	}
182	search := node.Name()
183	if strings.Contains(pattern, "*") {
184		search = node.path
185	}
186	re, err := regexp.Compile(prefix + pattern)
187	return err == nil && re.FindString(search) != ""
188}
189
190func (node *Node) sort(opts *Options) {
191	var fn SortFunc
192	switch {
193	case opts.ModSort:
194		fn = ModSort
195	case opts.CTimeSort:
196		fn = CTimeSort
197	case opts.DirSort:
198		fn = DirSort
199	case opts.VerSort:
200		fn = VerSort
201	case opts.SizeSort:
202		fn = SizeSort
203	case opts.NameSort:
204		fn = NameSort
205	default:
206		fn = NameSort // Default should be sorted, not unsorted.
207	}
208	if fn != nil {
209		if opts.ReverSort {
210			sort.Sort(sort.Reverse(ByFunc{node.nodes, fn}))
211		} else {
212			sort.Sort(ByFunc{node.nodes, fn})
213		}
214	}
215}
216
217// Path returns the Node's absolute path
218func (node *Node) Path() string {
219	return node.path
220}
221
222// Print nodes based on the given configuration.
223func (node *Node) Print(opts *Options) { node.print("", opts) }
224
225func dirRecursiveSize(opts *Options, node *Node) (size int64, err error) {
226	if opts.DeepLevel > 0 && node.depth >= opts.DeepLevel {
227		err = errors.New("Depth too high")
228	}
229
230	for _, nnode := range node.nodes {
231		if nnode.err != nil {
232			err = nnode.err
233			continue
234		}
235
236		if !nnode.IsDir() {
237			size += nnode.Size()
238		} else {
239			nsize, e := dirRecursiveSize(opts, nnode)
240			size += nsize
241			if e != nil {
242				err = e
243			}
244		}
245	}
246	return
247}
248
249func (node *Node) print(indent string, opts *Options) {
250	if node.err != nil {
251		err := node.err.Error()
252		if msgs := strings.Split(err, ": "); len(msgs) > 1 {
253			err = msgs[1]
254		}
255		fmt.Printf("%s [%s]\n", node.path, err)
256		return
257	}
258	if !node.IsDir() {
259		var props []string
260		ok, inode, device, uid, gid := getStat(node)
261		// inodes
262		if ok && opts.Inodes {
263			props = append(props, fmt.Sprintf("%d", inode))
264		}
265		// device
266		if ok && opts.Device {
267			props = append(props, fmt.Sprintf("%3d", device))
268		}
269		// Mode
270		if opts.FileMode {
271			props = append(props, node.Mode().String())
272		}
273		// Owner/Uid
274		if ok && opts.ShowUid {
275			uidStr := strconv.Itoa(int(uid))
276			if u, err := user.LookupId(uidStr); err != nil {
277				props = append(props, fmt.Sprintf("%-8s", uidStr))
278			} else {
279				props = append(props, fmt.Sprintf("%-8s", u.Username))
280			}
281		}
282		// Gorup/Gid
283		// TODO: support groupname
284		if ok && opts.ShowGid {
285			gidStr := strconv.Itoa(int(gid))
286			props = append(props, fmt.Sprintf("%-4s", gidStr))
287		}
288		// Size
289		if opts.ByteSize || opts.UnitSize {
290			var size string
291			if opts.UnitSize {
292				size = fmt.Sprintf("%4s", formatBytes(node.Size()))
293			} else {
294				size = fmt.Sprintf("%11d", node.Size())
295			}
296			props = append(props, size)
297		}
298		// Last modification
299		if opts.LastMod {
300			props = append(props, node.ModTime().Format("Jan 02 15:04"))
301		}
302		// Print properties
303		if len(props) > 0 {
304			fmt.Fprintf(opts.OutFile, "[%s]  ", strings.Join(props, " "))
305		}
306	} else {
307		var props []string
308		// Size
309		if opts.ByteSize || opts.UnitSize {
310			var size string
311			rsize, err := dirRecursiveSize(opts, node)
312			if err != nil && rsize <= 0 {
313				if opts.UnitSize {
314					size = "????"
315				} else {
316					size = "???????????"
317				}
318			} else if opts.UnitSize {
319				size = fmt.Sprintf("%4s", formatBytes(rsize))
320			} else {
321				size = fmt.Sprintf("%11d", rsize)
322			}
323			props = append(props, size)
324		}
325		// Print properties
326		if len(props) > 0 {
327			fmt.Fprintf(opts.OutFile, "[%s]  ", strings.Join(props, " "))
328		}
329	}
330	// name/path
331	var name string
332	if node.depth == 0 || opts.FullPath {
333		name = node.path
334	} else {
335		name = node.Name()
336	}
337	// Quotes
338	if opts.Quotes {
339		name = fmt.Sprintf("\"%s\"", name)
340	}
341	// Colorize
342	if opts.Colorize {
343		name = opts.color(node, name)
344	}
345	// IsSymlink
346	if node.Mode()&os.ModeSymlink == os.ModeSymlink {
347		vtarget, err := os.Readlink(node.path)
348		if err != nil {
349			vtarget = node.path
350		}
351		targetPath, err := filepath.EvalSymlinks(node.path)
352		if err != nil {
353			targetPath = vtarget
354		}
355		fi, err := opts.Fs.Stat(targetPath)
356		if opts.Colorize && fi != nil {
357			vtarget = opts.color(&Node{FileInfo: fi, path: vtarget}, vtarget)
358		}
359		name = fmt.Sprintf("%s -> %s", name, vtarget)
360		// Follow symbolic links like directories
361		if opts.FollowLink {
362			path, err := filepath.Abs(targetPath)
363			if err == nil && fi != nil && fi.IsDir() {
364				if _, ok := node.vpaths[filepath.Clean(path)]; !ok {
365					inf := &Node{FileInfo: fi, path: targetPath}
366					inf.vpaths = node.vpaths
367					inf.Visit(opts)
368					node.nodes = inf.nodes
369				} else {
370					name += " [recursive, not followed]"
371				}
372			}
373		}
374	}
375	// Print file details
376	// the main idea of the print logic came from here: github.com/campoy/tools/tree
377	fmt.Fprintln(opts.OutFile, name)
378	add := "│   "
379	for i, nnode := range node.nodes {
380		if opts.NoIndent {
381			add = ""
382		} else {
383			if i == len(node.nodes)-1 {
384				fmt.Fprintf(opts.OutFile, indent+"└── ")
385				add = "    "
386			} else {
387				fmt.Fprintf(opts.OutFile, indent+"├── ")
388			}
389		}
390		nnode.print(indent+add, opts)
391	}
392}
393
394const (
395	_        = iota // ignore first value by assigning to blank identifier
396	KB int64 = 1 << (10 * iota)
397	MB
398	GB
399	TB
400	PB
401	EB
402)
403
404// Convert bytes to human readable string. Like a 2 MB, 64.2 KB, 52 B
405func formatBytes(i int64) (result string) {
406	var n float64
407	sFmt, eFmt := "%.01f", ""
408	switch {
409	case i > EB:
410		eFmt = "E"
411		n = float64(i) / float64(EB)
412	case i > PB:
413		eFmt = "P"
414		n = float64(i) / float64(PB)
415	case i > TB:
416		eFmt = "T"
417		n = float64(i) / float64(TB)
418	case i > GB:
419		eFmt = "G"
420		n = float64(i) / float64(GB)
421	case i > MB:
422		eFmt = "M"
423		n = float64(i) / float64(MB)
424	case i > KB:
425		eFmt = "K"
426		n = float64(i) / float64(KB)
427	default:
428		sFmt = "%.0f"
429		n = float64(i)
430	}
431	if eFmt != "" && n >= 10 {
432		sFmt = "%.0f"
433	}
434	result = fmt.Sprintf(sFmt+eFmt, n)
435	result = strings.Trim(result, " ")
436	return
437}
438