From f09e2244da6f460d1786d7a3d668db9cebb5f790 Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <alxchk@gmail.com>
Date: Wed, 19 Oct 2016 23:14:56 +0300
Subject: [PATCH] Add mountinfo output for Linux to drives module

---
 pupy/modules/drives.py           |  56 ++++++++++++-----
 pupy/packages/linux/all/mount.py | 102 +++++++++++++++++++++++++++++++
 2 files changed, 143 insertions(+), 15 deletions(-)
 create mode 100644 pupy/packages/linux/all/mount.py

diff --git a/pupy/modules/drives.py b/pupy/modules/drives.py
index 0132b58a..5a3f36e2 100644
--- a/pupy/modules/drives.py
+++ b/pupy/modules/drives.py
@@ -1,22 +1,48 @@
-# -*- coding: UTF8 -*-
+# -*- coding: utf-8 -*-
 from pupylib.PupyModule import *
-# import ctypes
 from pupylib.utils.rpyc_utils import redirected_stdio
 
 __class_name__="Drives"
 
-@config(compat="windows", category="admin")
+@config(compat=[ 'linux', 'windows' ], category='admin')
 class Drives(PupyModule):
-	""" List valid drives in the system """
-	
-	dependencies=["win32api","win32com","pythoncom","winerror"]
-	
-	def init_argparse(self):
-		self.arg_parser = PupyArgumentParser(prog="drives", description=self.__doc__)
+    """ List valid drives in the system """
 
-	def run(self, args):
-		self.client.load_package("wmi")
-		self.client.load_package("pupwinutils.drives")
-		
-		with redirected_stdio(self.client.conn):
-			self.client.conn.modules['pupwinutils.drives'].list_drives()
\ No newline at end of file
+    dependencies={
+        'windows': [
+            'win32api', 'win32com', 'pythoncom',
+            'winerror', 'wmi', 'pupwinutils.drives'
+        ],
+        'linux': [ 'mount' ]
+    }
+
+    def init_argparse(self):
+        self.arg_parser = PupyArgumentParser(
+            prog="drives",
+            description=self.__doc__
+        )
+
+    def run(self, args):
+        if self.client.is_windows():
+            with redirected_stdio(self.client.conn):
+                self.client.conn.modules['pupwinutils.drives'].list_drives()
+
+        elif self.client.is_linux():
+            mountinfo = self.client.conn.modules['mount'].mounts()
+            for fstype in mountinfo.iterkeys():
+                if fstype in ('regular', 'dm'):
+                    continue
+
+                print '{}:'.format(fstype)
+                for info in mountinfo[fstype]:
+                    print info
+                print ''
+
+            for fstype in [ 'regular', 'dm' ]:
+                if not fstype in mountinfo:
+                    continue
+
+                print '{}: '.format(fstype)
+                for info in mountinfo[fstype]:
+                    print info
+                print ''
diff --git a/pupy/packages/linux/all/mount.py b/pupy/packages/linux/all/mount.py
new file mode 100644
index 00000000..8659bb07
--- /dev/null
+++ b/pupy/packages/linux/all/mount.py
@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+
+import os
+
+class MountInfo(object):
+    def __init__(self, line):
+        src, dst, fsname, options, _, _ = (
+            option.replace(r'\040', ' ') for option in line.split(' ')
+        )
+
+        self.src = src
+        self.dst = dst
+        self.fsname = fsname
+        self._options = options.split(',')
+
+        try:
+            vfsstat = os.statvfs(self.dst)
+            self.free = vfsstat.f_bfree*vfsstat.f_bsize
+            self.total = vfsstat.f_blocks*vfsstat.f_bsize
+            self.files = vfsstat.f_files
+            self.exception = None
+        except Exception as e:
+            self.exception = e.message
+            self.total = None
+            self.free = None
+            self.files = None
+
+        if self.fsname == 'tmpfs':
+            self.fstype = 'tmpfs'
+        elif self.fsname.startswith('cgroup'):
+            self.fstype = 'cgroup'
+        elif self.src.startswith('/dev/mapper'):
+            self.fstype = 'dm'
+        elif fsname.startswith('fuse.'):
+            self.fstype = 'fuse'
+        elif self.src == 'systemd-1' and self.fsname == 'autofs':
+            self.fstype = 'automount'
+        elif self.src == 'sunrpc':
+            self.fstype = 'rpc'
+        elif src == fsname:
+            self.fstype = 'kernel'
+        else:
+            self.fstype = 'block'
+
+    @property
+    def options(self):
+        return [
+            option for option in self._options if not any([
+                option.startswith(value) for value in (
+                    'relatime', 'fd', 'pgrp', 'timeout',
+                    'minproto', 'maxproto', 'direct', 'pipe_ino',
+                    'iocharset', 'codepage', 'lazytime', 'background_gc',
+                    'inline_data', 'discard', 'flush_merge', 'extent_cache',
+                    'mode', 'active_logs', 'commit', 'data', 'nr_inodes', 'size',
+                    'shortnames', 'utf8', 'errors'
+                )
+            ])
+        ]
+
+    def __repr__(self):
+        if self.fsname in ( 'tmpfs', 'cgroup', 'fuse', 'automount', 'rpc', 'kernel' ):
+            return '{} type={} options={}'.format(
+                self.dst, self.fstype, ','.join(self.options)
+            )
+        else:
+            Kb = 1024
+            Mb = 1024*Kb
+            Gb = 1024*Mb
+            Tb = 1024*Gb
+
+            if self.total > 0:
+                if self.free > Tb:
+                    free = '{}Tb'.format(self.free / Tb)
+                elif self.free > Gb:
+                    free = '{}Gb'.format(self.free / Gb)
+                elif self.free > Mb:
+                    free = '{}Mb'.format(self.free / Mb)
+                elif self.free > Kb:
+                    free = '{}Kb'.format(self.free / Kb)
+                else:
+                    free = '{}b'.format(self.free)
+
+                free = ' free={}%({})'.format(int(self.free/float(self.total)*100), free)
+            elif self.exception:
+                free = ' free=(error: {})'.format(self.exception)
+            else:
+                free = ' '
+
+            return '{} src={} fs={} options={}{}'.format(
+                self.dst, self.src, self.fsname, ','.join(self.options), free
+            )
+
+def mounts():
+    mountinfo = {}
+    with open('/proc/self/mounts', 'r') as mounts:
+        for line in mounts:
+            info = MountInfo(line)
+            if not info.fstype in mountinfo:
+                mountinfo[info.fstype] = [ info ]
+            else:
+                mountinfo[info.fstype].append(info)
+    return mountinfo