1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
|
#!/usr/bin/env python3
#
# Copyright (C) 2019-2021 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# 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/>.
import argparse
import os
import re
from datetime import datetime
from vyos.util import ask_yes_no
from vyos.util import call
from vyos.util import cmd
from vyos.util import DEVNULL
def list_disks():
disks = set()
with open('/proc/partitions') as partitions_file:
for line in partitions_file:
fields = line.strip().split()
if len(fields) == 4 and fields[3].isalpha() and fields[3] != 'name':
disks.add(fields[3])
return disks
def is_busy(disk: str):
"""Check if given disk device is busy by re-reading it's partition table"""
return call(f'blockdev --rereadpt /dev/{disk}', stderr=DEVNULL) != 0
def backup_partitions(disk: str):
"""Save sfdisk partitions output to a backup file"""
device_path = f'/dev/{disk}'
backup_ts = datetime.now().strftime('%Y%m%d-%H%M')
backup_file = f'/var/tmp/backup_{disk}.{backup_ts}'
call(f'sfdisk -d {device_path} > {backup_file}')
print(f'Partition table backup saved to {backup_file}')
def list_partitions(disk: str):
"""List partition numbers of a given disk"""
parts = set()
part_num_expr = re.compile(disk + '([0-9]+)')
with open('/proc/partitions') as partitions_file:
for line in partitions_file:
fields = line.strip().split()
if len(fields) == 4 and fields[3] != 'name' and part_num_expr.match(fields[3]):
part_idx = part_num_expr.match(fields[3]).group(1)
parts.add(int(part_idx))
return parts
def delete_partition(disk: str, partition_idx: int):
cmd(f'parted /dev/{disk} rm {partition_idx}')
def format_disk_like(target: str, proto: str):
cmd(f'sfdisk -d /dev/{proto} | sfdisk --force /dev/{target}')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
group = parser.add_argument_group()
group.add_argument('-t', '--target', type=str, required=True, help='Target device to format')
group.add_argument('-p', '--proto', type=str, required=True, help='Prototype device to use as reference')
args = parser.parse_args()
target_disk = args.target
eligible_target_disks = list_disks()
proto_disk = args.proto
eligible_proto_disks = eligible_target_disks.copy()
eligible_proto_disks.remove(target_disk)
if proto_disk == target_disk:
print('The two disk drives must be different.')
exit(1)
if not os.path.exists(f'/dev/{proto_disk}'):
print(f'Device /dev/{proto_disk} does not exist')
exit(1)
if not os.path.exists('/dev/' + target_disk):
print(f'Device /dev/{target_disk} does not exist')
exit(1)
if target_disk not in eligible_target_disks:
print(f'Device {target_disk} can not be formatted')
exit(1)
if proto_disk not in eligible_proto_disks:
print(f'Device {proto_disk} can not be used as a prototype for {target_disk}')
exit(1)
if is_busy(target_disk):
print(f'Disk device {target_disk} is busy, unable to format')
exit(1)
print(f'\nThis will re-format disk {target_disk} so that it has the same disk'
f'\npartion sizes and offsets as {proto_disk}. This will not copy'
f'\ndata from {proto_disk} to {target_disk}. But this will erase all'
f'\ndata on {target_disk}.\n')
if not ask_yes_no('Do you wish to proceed?'):
print(f'Disk drive {target_disk} will not be re-formated')
exit(0)
print(f'Re-formating disk drive {target_disk}...')
print('Making backup copy of partitions...')
backup_partitions(target_disk)
print('Deleting old partitions...')
for p in list_partitions(target_disk):
delete_partition(disk=target_disk, partition_idx=p)
print(f'Creating new partitions on {target_disk} based on {proto_disk}...')
format_disk_like(target=target_disk, proto=proto_disk)
print('Done!')
|
