mirror of https://github.com/deadc0de6/catcli
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
211 lines
7.4 KiB
Python
211 lines
7.4 KiB
Python
"""
|
|
author: deadc0de6 (https://github.com/deadc0de6)
|
|
Copyright (c) 2017, deadc0de6
|
|
|
|
Class that represents a node in the catalog tree
|
|
"""
|
|
|
|
import os
|
|
import anytree
|
|
import psutil
|
|
import time
|
|
|
|
# local imports
|
|
import catcli.utils as utils
|
|
from catcli.logger import Logger
|
|
|
|
'''
|
|
There are 4 types of node:
|
|
* "top" node representing the top node (generic node)
|
|
* "storage" node representing a storage
|
|
* "dir" node representing a directory
|
|
* "file" node representing a file
|
|
'''
|
|
|
|
|
|
class Noder:
|
|
|
|
TOPNAME = 'top'
|
|
TYPE_TOP = 'top' # tip top ;-)
|
|
TYPE_FILE = 'file'
|
|
TYPE_DIR = 'dir'
|
|
TYPE_STORAGE = 'storage'
|
|
|
|
def __init__(self, verbose=False):
|
|
self.hash = True
|
|
self.verbose = verbose
|
|
|
|
def set_hashing(self, val):
|
|
self.hash = val
|
|
|
|
def get_storage_names(self, top):
|
|
''' return a list of all storage names '''
|
|
return [x.name for x in list(top.children)]
|
|
|
|
###############################################################
|
|
# node creationg
|
|
###############################################################
|
|
def new_top_node(self):
|
|
''' create a new top node'''
|
|
return anytree.AnyNode(name=self.TOPNAME, type=self.TYPE_TOP)
|
|
|
|
def file_node(self, name, path, parent, storagepath):
|
|
''' create a new node representing a file '''
|
|
path = os.path.abspath(path)
|
|
st = os.lstat(path)
|
|
md5 = None
|
|
if self.hash:
|
|
md5 = utils.md5sum(path)
|
|
relpath = os.path.join(os.path.basename(storagepath),
|
|
os.path.relpath(path, start=storagepath))
|
|
return self._node(name, self.TYPE_FILE, relpath, parent,
|
|
size=st.st_size, md5=md5)
|
|
|
|
def dir_node(self, name, path, parent, storagepath):
|
|
''' create a new node representing a directory '''
|
|
path = os.path.abspath(path)
|
|
relpath = os.path.relpath(path, start=storagepath)
|
|
return self._node(name, self.TYPE_DIR, relpath, parent)
|
|
|
|
def storage_node(self, name, path, parent, attr=None):
|
|
''' create a new node representing a storage '''
|
|
path = os.path.abspath(path)
|
|
free = psutil.disk_usage(path).free
|
|
total = psutil.disk_usage(path).total
|
|
epoch = time.time()
|
|
return anytree.AnyNode(name=name, type=self.TYPE_STORAGE, free=free,
|
|
total=total, parent=parent, attr=attr, ts=epoch)
|
|
|
|
def _node(self, name, type, relpath, parent, size=None, md5=None):
|
|
''' generic node creation '''
|
|
return anytree.AnyNode(name=name, type=type, relpath=relpath,
|
|
parent=parent, size=size, md5=md5)
|
|
|
|
###############################################################
|
|
# printing
|
|
###############################################################
|
|
def _print_node(self, node, pre='', withpath=False, withdepth=False):
|
|
''' print a node '''
|
|
if node.type == self.TYPE_TOP:
|
|
Logger.out('{}{}'.format(pre, node.name))
|
|
elif node.type == self.TYPE_FILE:
|
|
name = node.name
|
|
if withpath:
|
|
name = node.relpath
|
|
attr = ''
|
|
if node.md5:
|
|
attr = ', md5:{}'.format(node.md5)
|
|
compl = 'size:{}{}'.format(utils.human(node.size), attr)
|
|
Logger.file(pre, name, compl)
|
|
elif node.type == self.TYPE_DIR:
|
|
name = node.name
|
|
if withpath:
|
|
name = node.relpath
|
|
depth = ''
|
|
if withdepth:
|
|
depth = len(node.children)
|
|
attr = None
|
|
if node.size:
|
|
attr = [['totsize', utils.human(node.size)]]
|
|
Logger.dir(pre, name, depth=depth, attr=attr)
|
|
elif node.type == self.TYPE_STORAGE:
|
|
hf = utils.human(node.free)
|
|
ht = utils.human(node.total)
|
|
name = '{} (free:{}, total:{})'.format(node.name, hf, ht)
|
|
Logger.storage(pre, name, node.attr)
|
|
else:
|
|
Logger.err('Weird node encountered: {}'.format(node))
|
|
# Logger.out('{}{}'.format(pre, node.name))
|
|
|
|
def print_tree(self, node, style=anytree.ContRoundStyle()):
|
|
''' print the tree similar to unix tool "tree" '''
|
|
for pre, fill, node in anytree.RenderTree(node, childiter=self.sort_tree):
|
|
self._print_node(node, pre=pre, withdepth=True)
|
|
|
|
###############################################################
|
|
# searching
|
|
###############################################################
|
|
def find_name(self, root, key, script=False):
|
|
''' find files based on their names '''
|
|
if self.verbose:
|
|
Logger.info('searching for \"{}\"'.format(key))
|
|
self.term = key
|
|
found = anytree.findall(root, filter_=self._find_name)
|
|
paths = []
|
|
for f in found:
|
|
if f.type == self.TYPE_STORAGE:
|
|
# ignore storage nodes
|
|
continue
|
|
self._print_node(f, withpath=True, withdepth=True)
|
|
paths.append(f.relpath)
|
|
if script:
|
|
tmp = ['${source}/'+x for x in paths]
|
|
cmd = 'op=file; source=/media/mnt; $op {}'.format(' '.join(tmp))
|
|
Logger.info(cmd)
|
|
return found
|
|
|
|
def _find_name(self, node):
|
|
''' callback for finding files '''
|
|
if self.term.lower() in node.name.lower():
|
|
return True
|
|
return False
|
|
|
|
###############################################################
|
|
# climbing
|
|
###############################################################
|
|
def walk(self, root, path, rec=False):
|
|
''' walk the tree for ls based on names '''
|
|
if self.verbose:
|
|
Logger.info('walking path: \"{}\"'.format(path))
|
|
r = anytree.resolver.Resolver('name')
|
|
found = []
|
|
try:
|
|
found = r.glob(root, path)
|
|
if len(found) < 1:
|
|
return []
|
|
if rec:
|
|
self.print_tree(found[0].parent)
|
|
return
|
|
found = sorted(found, key=self.sort_walk)
|
|
self._print_node(found[0].parent,
|
|
withpath=False, withdepth=True)
|
|
for f in found:
|
|
self._print_node(f, withpath=False, pre='- ', withdepth=True)
|
|
except anytree.resolver.ChildResolverError:
|
|
pass
|
|
return found
|
|
|
|
###############################################################
|
|
# diverse
|
|
###############################################################
|
|
def sort_tree(self, items):
|
|
''' sorting a list of items '''
|
|
return sorted(items, key=self.sort_walk)
|
|
|
|
def sort_walk(self, n):
|
|
''' for sorting a node '''
|
|
return (n.type, n.name.lstrip('\.').lower())
|
|
|
|
def to_dot(self, node, path='tree.dot'):
|
|
''' export to dot for graphing '''
|
|
anytree.exporter.DotExporter(node).to_dotfile(path)
|
|
Logger.info('dot file created under \"{}\"'.format(path))
|
|
return 'dot {} -T png -o /tmp/tree.png'.format(path)
|
|
|
|
def rec_size(self, node):
|
|
''' recursively traverse tree and store dir size '''
|
|
if self.verbose:
|
|
Logger.info('getting folder size recursively')
|
|
if node.type == self.TYPE_FILE:
|
|
return node.size
|
|
size = 0
|
|
for i in node.children:
|
|
if node.type == self.TYPE_DIR:
|
|
size += self.rec_size(i)
|
|
if node.type == self.TYPE_STORAGE:
|
|
self.rec_size(i)
|
|
else:
|
|
continue
|
|
node.size = size
|
|
return size
|