btrfs-rm

#!/usr/bin/env python3
import argparse
import hashlib
import io
import logging
import os
import subprocess
import sys

DRY_RUN=True
DEBUG_LEVEL=30

class FileSystemEntry:
    source  = ""
    target  = ""
    mtype   = ""
    options = ""
    freq    = ""
    passno  = ""

    def __str__(self):
        return self.source + " " + self.target + " " + self.mtype + " " + self.options

class Mount(object):

    def __init__(self, proc_mounts_file='/proc/mounts'):
        self.entries = self.load(proc_mounts_file)

    def read_file_into_array(self,filename):
        file = open(filename, 'r')
        array = file.readlines()
        file.close()
        return array

    def parse(self,array):
        entries = []
        for line in array:
            cols = line.split()
            fse  = FileSystemEntry()
            fse.source  = cols[0]
            fse.target  = cols[1]
            fse.mtype   = cols[2]
            fse.options = cols[3]
            fse.freq    = cols[4]
            fse.passno  = cols[5]
            entries.append(fse)
        return entries

    def load(self, proc_mounts_file):
        array   = self.read_file_into_array(proc_mounts_file)
        entries = self.parse(array)
        return entries

    # Filters all virtual filesystems out (they do not start with '/')
    def filter_by_path(self,entries):
        result = []	
        for o in self.entries:
            if o.source.startswith(os.sep):
               result.append(o)
        return result

    def longest_match(self, path, entries):
        max_len = 0
        result = None
        for o in entries:
            l = len(o.target)
            if path.startswith(o.target) and l > max_len:
                    result = o
                    max_len = l
        return result

    def get_subvol(self, entry):
        options = entry.options.split(",")
        for option in options:
            if option.startswith("subvol="):
                return option[len("subvol="):]
        return None

    def find_root(self, source, entries):
        for o in entries:
            if o.source == source and self.get_subvol(o) == os.sep:
                return o
        return None

    def resolve(self, path, entries):
        entry = self.longest_match(path,entries)
        if entry is None:
            return path
        if entry.target == os.sep:
            return path

        subvol = self.get_subvol(entry)

        if subvol is None:
            return path

        if subvol != os.sep:
            postfix = path[len(entry.target):]
            _path   = subvol + postfix
            if _path != entry.target:
                path    = self.resolve(_path,entries)

        root = self.find_root(entry.source,entries)

        if root is None:
            return path

        if path == root.target or path.startswith(root.target + os.sep):
            return path

        return (root.target + path)

    def resolve_bind(self, path):
        result = self.filter_by_path(self.entries)
        result = self.resolve(path,result)
        return result

class MountRecorder(Mount):
    def load(self, proc_mounts_file):
        array   = self.read_file_into_array(proc_mounts_file)
        content = "".join(array) 
        text_file = open("rec_mounts", "w")
        text_file.write(content)
        text_file.close()

        entries = self.parse(array)

        return entries


class ShellCmd(object):

    def subprocess(self, args):
        proc = subprocess.Popen(args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
        return proc

    def read_stdout(self, stream):
        array=[] 
        for line in stream:
            item=line.rstrip("\n")
            array.append(item)
        return array

    # The exit_code can be access from external then
    def execute(self, args):
        proc        = self.subprocess(args)

        # Python2 vs Python3
        if sys.version_info >= (3,0):
            input_stream = io.TextIOWrapper(proc.stdout, encoding='utf-8')
        else:
            input_stream = proc.stdout

        array       = self.read_stdout(input_stream)
        exit_code   = proc.wait()
        return ( exit_code, array )

class ShellCmdRecorder(ShellCmd):
    def execute(self, args):
        exit_code, array = super(ShellCmdRecorder,self).execute(args)

        cmdline = " ".join(args)
        filename_sha1 = hashlib.sha1(cmdline.encode("utf-8")).hexdigest()
        content = "\n".join(array)

        text_file = open("rec_"+filename_sha1, "w")
        text_file.write(cmdline + "\n")
        text_file.write(str(exit_code) + "\n")
        text_file.write(content + "\n")
        text_file.close()

        return (exit_code,array)


class BTRFS(object):

  def __init__(self, shellCmd, mount):
    self.cmd   = shellCmd
    self.mount = mount

  def show(self, btrfs_dir, verbose=False):
    args = ['btrfs','subvolume','show', btrfs_dir]
    logger.debug(" ".join(args))
    return self.cmd.execute(args)

  def find_btrfs(self, btrfs_dir):
    while True:
        exit_code,array = self.show(btrfs_dir)
        if btrfs_dir == os.sep:
            break
        if exit_code != 1:
            break
        # Move one level up
        btrfs_dir = os.path.dirname(btrfs_dir)
    if exit_code != 0:
        return (None,None)
    if len(array) == 1:
        btrfs_name = ""
    else:
        btrfs_name = array[1].split()[1]
    return (btrfs_dir, btrfs_name)

  def subvolume_list_raw(self,btrfs_dir):
    args = ['btrfs','subvolume', 'list', '-o', btrfs_dir ]
    logger.debug(" ".join(args))
    return self.cmd.execute(args)

  def subvolume_list_extract_dirs(self, array):
    result = []
    for line in array:
      columns     = line.split()
      last_column = len(columns) - 1
      dir_path    = columns[last_column]
      result.append(dir_path)
    return result

  def subvolume_list_filter(self,target_dir,btrfs_dir,btrfs_name, source2, array):
    array   = self.subvolume_list_extract_dirs(array)
    if btrfs_name == "<FS_TREE>":
        filter1 = source2[1:] 
    else:
        filter1 = btrfs_name + os.sep
   
    length  = len(filter1)
    array2  = []

    for item in array:
        # filter for the btrfs_name prefix name
        if item.startswith(filter1):
            item2 = item[length:]
            if btrfs_name == "<FS_TREE>":
                item3 = btrfs_dir + os.sep +  filter1 + item2
            else:
                item3 = btrfs_dir + item2
            if target_dir == item3 or item3.startswith(target_dir + os.sep):
                array2.append(item3)
    return array2


  def subvolume_list(self,btrfs_dir, btrfs_name, source2, target_dir):
        exit_code, array= self.subvolume_list_raw(btrfs_dir)
        if exit_code != 0:
            return []
        subvolumes = self.subvolume_list_filter(target_dir,btrfs_dir,btrfs_name,source2,array)
        return subvolumes

  def delete(self,dirname):
    args = ['btrfs','subvolume', 'delete', dirname]
    logger.info(" ".join(args))
    if DRY_RUN == True:
        return 

    self.cmd.execute(args)
    logger.debug("delete finished")
    # FIXME return code?

  def findmnt(self):
    # findmnt
    # -c Canonicalize all printed paths
    # -n No Headings
    # -D Imitate the output of df
    args = ['findmnt','-c','-n','-D']
    logger.debug(" ".join(args))
    return self.cmd.execute(args)


  def parse_btrfs_mount(self,source):
    pos = source.find("[")
    if pos == -1:
        return [ source, ''  ]
    pos_end = source.find("]", pos)
    if pos_end == -1:
        return [ source, '' ]

    source1 = source[0:pos]
    source2 = source[pos+1:pos_end]

    return [ source1, source2 ]


  def filter_findmnt_target(self,lines):
    result = []

    for line in lines:
      columns     = line.split()
      last_column = len(columns) - 1
      source      = columns[0]          
      target      = columns[last_column] 

      # source can only be device path starting with a '/', everthing else is shm, tmpfs, etc.
      if not source.startswith(os.sep):
          continue
      source1,source2      = self.parse_btrfs_mount(source)
      result.append([source1,source2,target])

    return result

  # In the end this method wants to match the closed "mountpoint" coming from a bottom up to the tree
  # So if the pathname is longer than any target mount points, we remove components of the dirname, to find
  # a final "exact" match. This means we have then a (partial) match of the dirname path, which goes to the
  # level where the mountpoint is

  # Test case 1: dirname is shorter  --> match root '/'
  # Test case 2: dirname is longer --> match mount point which matches first when coming from bottom (longest path) 
  # so there is someimes an exact match, somethings not

  def match_exact(self,name,lines):
    i = 0
    for line in lines:
        source1,source2,target = line
        if target == name:
            return i
        i=i+1
    return -1

  def match_dirname(self,dirname,lines):

    components = dirname.split(os.sep)
    length     = len(components)

    for i in reversed(range (0,length)):

        if    i == 0: name = os.sep
        else:         name = os.sep.join(components[0:i+1])

        pos = self.match_exact(name,lines)

        if pos > -1:
            result = lines[pos]
            break

    # There has to be always a match, otherwise this dirname was no absolute path
    return result

  def find_longest_entry(self,source1,filtered_lines):
    result = ""
    for line in filtered_lines:
        _source1,_source2,_target = line
        if _source1 == source1 and _source2 == '':
            return _target
    return ""

  def find_match(self,dirname):
    exit_code, lines = self.findmnt()
    if not exit_code == 0:
      return []
    filtered_lines  = self.filter_findmnt_target(lines)

    result                  = self.match_dirname(dirname,filtered_lines)
    source1,source2,target  = result
    return [ source1, source2, target, dirname ]

  def rm_recursive(self,dirname):
    args = [ 'rm', '-rf', dirname]
    logger.info(" ".join(args))
    if DRY_RUN == True:
     return 

    proc = subprocess.Popen(args,stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    for line in proc.stdout:
      print(line.strip())
    exit_code = proc.wait()

  def get_path(self):
    try:
        path_var   = os.environ['PATH']
        path_array = path_var.split(os.pathsep)
        return path_array
    except KeyError:
        return []

  def find_program(self,name,path_array):
    for path in path_array:
        if not os.path.isdir(path):
            continue
        filepath = os.path.join(path,name) 
        if os.path.isfile(filepath):
            return True
    return False

  def find_programs(self,required_programs,path_array):
    not_found_list = []
    for name in required_programs:
        found = self.find_program(name,path_array)
        if not found:
            not_found_list.append(name)
    return not_found_list

  def self_check(self,required_programs,path_array):
    not_found_list = self.find_programs(required_programs,path_array)

    if len(not_found_list) > 0 :
        msg = "Necessary program not found in PATH: " +  ",".join(not_found_list)
        logger.critical(msg)
        return False
    return True

  # Only Dirs as input (no files, no links etc.)
  def get_subvolumes(self, dirname):
    source1,source2,target,dirname = self.find_match(dirname)

    btrfs_dir, btrfs_name = self.find_btrfs(target)

    if btrfs_dir is None:
        return []

    logger.debug("btrfs_dir:  "+btrfs_dir)
    logger.debug("btrfs_name: "+btrfs_name)
    logger.debug("target:     "+target)

    subvolumes = self.subvolume_list(btrfs_dir, btrfs_name, source2, dirname)

    logger.debug("Subdirs")
    logger.debug("\n".join(subvolumes))

    return subvolumes

  # FIXME Only Dirs as input (no files, no links etc.)
  def process_dir(self, dirname):
    subvolumes = self.get_subvolumes(dirname)

    for subvolume in subvolumes:
      self.delete(subvolume)
    self.rm_recursive(dirname)

  def run(self,path_list):
    path_array        = self.get_path()
    required_programs = [ 'rm', 'btrfs', 'findmnt' ]
    if not self.self_check(required_programs,path_array):
        exit(1)


    for path in path_list:
        if os.path.isdir(path):
            abspath   = os.path.abspath(path)
            bind_path = self.mount.resolve_bind(abspath)
            self.process_dir(bind_path)
        else:
           self.rm_recursive(path)


def save_cmdline(args):
        content = "\n".join(args)
        text_file = open("rec_cmdline", "w")
        text_file.write(content+"\n")
        text_file.close()


if __name__ == '__main__':
        parser = argparse.ArgumentParser("btrfs-rm")
        parser.add_argument('directories', metavar='directory', nargs='+',
                    help='directory which contains btrfs snapshots ')
        parser.add_argument('-t', '--test',  dest='dry_run', action='store_true', default=False,  help='dry-run without actually deleting stuff')
        parser.add_argument('-d', '--debug', metavar='level', dest='debug_level', type=int, default=logging.WARN,  help='set debug level')
        parser.add_argument('-r', '--record',  dest='record', action='store_true', default=False,  help='record shell cmd results for analysis')

        args = parser.parse_args()
        DEBUG_LEVEL=args.debug_level
        DRY_RUN=args.dry_run
        if DRY_RUN and DEBUG_LEVEL > 20:
            DEBUG_LEVEL=20

        # after arg parse, so that somebody could at least use the -h option
        if os.geteuid() != 0:
            exit("You need to have root privileges to run this script.\nPlease try again, this time using 'sudo'.")
        # Create logger
        FORMAT = '%(message)s'
        logging.basicConfig(format=FORMAT)
        logger = logging.getLogger('btrfs-rm')
        logger.setLevel(DEBUG_LEVEL)

        if args.record:
            shellCmd = ShellCmdRecorder()
            mount    = MountRecorder()
            save_cmdline(sys.argv)
        else:
            shellCmd = ShellCmd()
            mount    = Mount()

        btrfs = BTRFS(shellCmd,mount)
        btrfs.run(args.directories)