format_disk.py 11.6 KB
Newer Older
1 2 3 4 5 6
#!/usr/bin/env python
# encoding: utf-8
from __future__ import print_function
from __future__ import division
from subprocess import CalledProcessError

7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29
''' 
/**
 * @file format_disk.py
 * @brief Prepare and partition new disk for fast recording. This script creates two partitions on a disk: 
 * one is formatted to ext4 and the other is left unformatted for fast recording from camogm.
 * @copyright Copyright (C) 2017 Elphel Inc.
 * @author Mikhail Karpenko <mikhail@elphel.com>
 * @deffield updated: 
 *
 * @par <b>License</b>:
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
30
'''
31

32
__author__ = "Elphel"
33
__copyright__ = "Copyright 2017 Elphel Inc."
34 35 36 37 38 39 40 41 42 43 44 45
__license__ = "GPL"
__version__ = "3.0+"
__maintainer__ = "Mikhail Karpenko"
__email__ = "mikhail@elphel.com"
__status__ = "Development"

import os
import re
import sys
import stat
import argparse
import subprocess
46
import time
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109

# use this % of total disk space for system partition
SYS_PARTITION_RATIO = 5

class ErrCodes(object):
    OK = 0
    WRONG_DISK = 1
    WRONG_PATH = 2
    NOT_DISK = 3
    NO_TOOLS = 4
    NO_PERMISSIONS = 5
    PART_FAILURE = 6
    def __init__(self, code = OK):
        """
        Prepare strings with error code description.
        """
        # the length of this list must match the number of error code defined as clas attributes
        self.err_str = ["Operation finished successfully", 
                        "The disk specified is already partitioned",
                        "Path to disk provided on the command line is invalid",
                        "The path provided is a partition, not a disk",
                        "One of the command-line utilities required for this script is not found",
                        "This scrip requires root permissions",
                        "Partitioning finished unsuccessfully"]
        self._err_code = code
    
    def err2str(self):
        """
        Convert error code to string description.
        Return: string containing error description
        """
        if self._err_code < len(self.err_str):
            ret = self.err_str[self._err_code]
        else:
            ret = "No description for this error code"
        return ret
    
    @property
    def err_code(self):
        return self._err_code

    @err_code.setter
    def err_code(self, val):
        if val < len(self.err_str):
            self._err_code = val

def check_prerequisites():
    """
    Check all tools required for disk partitioning.
    Return: emtry string if all tools are found and the name of a missing program otherwise
    """
    ret_str = ""
    check_tools = [['parted', '-v'],
                   ['mkfs.ext4', '-V']]
    # make it silent
    with open('/dev/null', 'w') as devnull:
        for tool in check_tools:
            try:
                subprocess.check_call(tool, stdout = devnull, stderr = devnull)
            except:
                ret_str = tool[0]
    return ret_str

110
def find_disks(partitioned = False):
111
    """
112 113 114
    Find all attached and, by default, unpartitioned SCSI disks. If a key is specified
    then a list of all attached disks is is returned.
    @param partitioned: include partitioned disks
115 116 117 118 119 120 121 122 123 124
    Return: a list containing paths to disks
    """
    dlist = []
    try:
        partitions = subprocess.check_output(['cat', '/proc/partitions'])
        # the first two elemets of the list are table header and empty line delimiter, skip them
        for partition in partitions.splitlines()[2:]:
            dev = re.search(' +(sd[a-z]$)', partition)
            if dev:
                dev_path = '/dev/{0}'.format(dev.group(1))
125 126 127 128 129
                if not partitioned:
                    plist = find_partitions(dev_path)
                    if not plist:
                        dlist.append(dev_path)
                else:
130 131 132 133 134 135 136 137 138 139 140 141 142 143 144
                    dlist.append(dev_path)
    except:
        # something went wrong, clear list to prevent accidental data loss
        del dlist[:]
    return dlist

def find_partitions(dev_path):
    """
    Find all partitions (if any) on a disk.
    @param dev_path: path to device
    Return: a list of full paths to partitions or empty list in case no partitions found on disk
    """
    plist = []
    try:
        partitions = subprocess.check_output(['cat', '/proc/partitions'])
145
        search_str = '([0-9]+) +({0}[0-9]+$)'.format(dev_path.rpartition('/')[-1])
146 147 148 149
        # the first two elemets of the list are table header and empty line delimiter, skip them
        for partition in partitions.splitlines()[2:]:
            dev = re.search(search_str, partition)
            if dev:
150
                plist.append('/dev/{0} ({1:.1f} GB)'.format(dev.group(2), int(dev.group(1)) / 1000000))
151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184
    except:
        # something went wrong, clear list to prevent accidental data loss
        del plist[:]
    return plist

def is_partition(dev_path):
    """
    Check if the path specified corresponds to partition and not to disk.
    @param dev_path: path to device
    Return: boolean value indicating if the path provided is a partition.
    """
    # disk path should end with a character only
    disk = re.search('sd[a-z]$', dev_path)
    if disk:
        ret = False
    else:
        ret = True
    return ret
    
def get_disk_size(dev_path):
    """
    Get the size of disk specified.
    @param dev_path: path to device
    Return: disk size in GB
    """
    try:
        parted_print = subprocess.check_output(['parted', '-m', dev_path, 'unit', 'GB', 'print'])
        fields = parted_print.split(':')
        sz = fields[1]
        disk_size = int(sz[:-2])
    except:
        disk_size = 0
    return disk_size

185
def partition_disk(dev_path, sys_size, disk_size, dry_run = True, force = False):
186 187 188 189 190 191 192 193 194 195 196 197
    """
    Create partitions on disk and format system partition.
    @param dev_path: path to device
    @param sys_size: the size of system partition in GB
    @param disk_size: total disk size in GB
    Return: empty string in case of success or error message indicating the result of partitioning
    """
    try:
        if not dry_run:
            # create system partition
            start = 0
            end = sys_size
198 199 200
            subprocess.check_output(['parted', '-s', dev_path, 'unit', 'GB',
                                     'mklabel', 'msdos',
                                     'mkpart', 'primary', str(start), str(end)], stderr = subprocess.STDOUT)
201 202 203
            # create raw partition
            start = sys_size
            end = disk_size
204 205
            subprocess.check_output(['parted', '-s', dev_path, 'unit', 'GB',
                                     'mkpart', 'primary', str(start), str(end)], stderr = subprocess.STDOUT)
206 207
            # make file system on first partition; delay to let the changes propagate to the system
            time.sleep(2)
208
            partition = dev_path + '1'
209
            if force:
210
                cmd_str = ['mkfs.ext4', '-FF', partition]
211 212 213 214 215 216 217 218
                # if system partition contained a file system then it will be mounted right after partitioning
                # check this situation and unmount partition
                mounted = subprocess.check_output(['mount'])
                for item in mounted.splitlines():
                    mount_point = re.search('^{0}'.format(partition), item)
                    if mount_point:
                        subprocess.check_output(['umount', partition])
            else:
219 220
                cmd_str = ['mkfs.ext4', partition]
            subprocess.check_output(cmd_str, stderr = subprocess.STDOUT)
221 222
        ret_str = ""
    except subprocess.CalledProcessError as e:
223
        ret_str = e.output
224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240
    except OSError as e:
        ret_str = e.strerror
    return ret_str

if __name__ == "__main__":
    ret_str = check_prerequisites()
    if ret_str != "":
        ret_code = ErrCodes(ErrCodes.NO_TOOLS)
        print("{0}: {1}".format(ret_code.err2str(), ret_str))
        sys.exit(ret_code.err_code)
    if os.geteuid() != 0:
        ret_code = ErrCodes(ErrCodes.NO_PERMISSIONS)
        print(ret_code.err2str())
        sys.exit(ret_code.err_code)

    parser = argparse.ArgumentParser(description = "Prepare and partition new disk for fast recording from camogm")
    parser.add_argument('disk_path', nargs = '?', help = "path to a disk which should be partitioned, e.g /dev/sda")
241 242
    parser.add_argument('-l', '--list', action = 'store_true', help = "list attached disk(s) suitable for partitioning along " + 
    "with their totals sizes and possible system partition sizes separated by colon")
243 244
    parser.add_argument('-e', '--errno', nargs = 1, type = int, help = "convert error number returned by the script to error message")
    parser.add_argument('-d', '--dry_run', action = 'store_true', help = "execute the script but do not actually create partitions")
245
    parser.add_argument('-f', '--force', action = 'store_true', help = "force 'mkfs' to create a file system")
246
    parser.add_argument('-p', '--partitions', action = 'store_true', help = "list partitions and their sizes separated by colon")
247 248 249 250 251
    args = parser.parse_args()

    if args.list:
        disks = find_disks()
        for disk in disks:
252 253 254 255 256 257
            total_size = get_disk_size(disk)
            if total_size > 0:
                sys_size = total_size * (SYS_PARTITION_RATIO / 100)
            else:
                sys_size = 0
            print('{0}:{1} GB:{2} GB'.format(disk, total_size, sys_size))
258 259 260 261 262 263
    elif args.partitions:
        all_partitions = []
        dlist = find_disks(partitioned = True)
        for disk in dlist:
            all_partitions += find_partitions(disk)
        print(':'.join(all_partitions))
264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294
    elif args.errno:
        ret = ErrCodes(args.errno[0])
        print(ret.err2str())
    elif args.disk_path:
        disk_path = ""
        ret_code = ErrCodes()
        if os.path.exists(args.disk_path):
            mode = os.stat(args.disk_path).st_mode
            if stat.S_ISBLK(mode):
                if not is_partition(args.disk_path):
                    disk_path = args.disk_path
                    plist = find_partitions(disk_path)
                    if not plist:
                        # OK, disk is not partitioned and we can proceed
                        ret_code.err_code = ErrCodes.OK
                    else:
                        # stop, disk is already partitioned
                        ret_code.err_code = ErrCodes.WRONG_DISK
                else:
                    ret_code.err_code = ErrCodes.NOT_DISK
            else:
                ret_code.err_code = ErrCodes.WRONG_PATH
        else:
            ret_code.err_code = ErrCodes.WRONG_PATH
        if ret_code.err_code != ErrCodes.OK:
            print(ret_code.err2str())
            sys.exit(ret_code.err_code)
            
        total_size = get_disk_size(disk_path)
        if total_size > 0:
            sys_size = total_size * (SYS_PARTITION_RATIO / 100)
295 296 297 298 299
            if args.force:
                force = args.force
            else:
                force = False
            ret_str = partition_disk(disk_path, sys_size, total_size, args.dry_run, force)
300 301
            if ret_str:
                ret_code = ErrCodes(ErrCodes.PART_FAILURE)
302
                print('{0}: {1}'.format(ret_code.err2str(), ret_str))
303 304 305
                sys.exit(ret_code.err_code)
    else:
        parser.print_help()