Statistics
| Branch: | Tag: | Revision:

root / snf-image-helper / disklabel.py @ 83d0c566

History | View | Annotate | Download (19.8 kB)

1 de34e7aa Nikos Skalkotos
#!/usr/bin/env python
2 1bda0902 Nikos Skalkotos
#
3 de34e7aa Nikos Skalkotos
# -*- coding: utf-8 -*-
4 de34e7aa Nikos Skalkotos
#
5 de34e7aa Nikos Skalkotos
# Copyright (C) 2013 GRNET S.A.
6 de34e7aa Nikos Skalkotos
#
7 de34e7aa Nikos Skalkotos
# This program is free software; you can redistribute it and/or modify
8 de34e7aa Nikos Skalkotos
# it under the terms of the GNU General Public License as published by
9 de34e7aa Nikos Skalkotos
# the Free Software Foundation; either version 2 of the License, or
10 de34e7aa Nikos Skalkotos
# (at your option) any later version.
11 de34e7aa Nikos Skalkotos
#
12 de34e7aa Nikos Skalkotos
# This program is distributed in the hope that it will be useful, but
13 de34e7aa Nikos Skalkotos
# WITHOUT ANY WARRANTY; without even the implied warranty of
14 de34e7aa Nikos Skalkotos
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15 de34e7aa Nikos Skalkotos
# General Public License for more details.
16 de34e7aa Nikos Skalkotos
#
17 de34e7aa Nikos Skalkotos
# You should have received a copy of the GNU General Public License
18 de34e7aa Nikos Skalkotos
# along with this program; if not, write to the Free Software
19 de34e7aa Nikos Skalkotos
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
20 de34e7aa Nikos Skalkotos
# 02110-1301, USA.
21 de34e7aa Nikos Skalkotos
22 de34e7aa Nikos Skalkotos
"""This module provides the code for handling OpenBSD disklabels"""
23 de34e7aa Nikos Skalkotos
24 de34e7aa Nikos Skalkotos
import struct
25 de34e7aa Nikos Skalkotos
import sys
26 de34e7aa Nikos Skalkotos
import cStringIO
27 de34e7aa Nikos Skalkotos
import optparse
28 de34e7aa Nikos Skalkotos
29 de34e7aa Nikos Skalkotos
from collections import namedtuple
30 de34e7aa Nikos Skalkotos
31 de34e7aa Nikos Skalkotos
BLOCKSIZE = 512
32 de34e7aa Nikos Skalkotos
33 de34e7aa Nikos Skalkotos
LABELSECTOR = 1
34 de34e7aa Nikos Skalkotos
LABELOFFSET = 0
35 de34e7aa Nikos Skalkotos
36 de34e7aa Nikos Skalkotos
BBSIZE = 8192  # size of boot area with label
37 de34e7aa Nikos Skalkotos
SBSIZE = 8192  # max size of fs superblock
38 de34e7aa Nikos Skalkotos
39 de34e7aa Nikos Skalkotos
DISKMAGIC = 0x82564557
40 de34e7aa Nikos Skalkotos
41 de34e7aa Nikos Skalkotos
42 de34e7aa Nikos Skalkotos
class MBR(object):
43 de34e7aa Nikos Skalkotos
    """Represents a Master Boot Record."""
44 de34e7aa Nikos Skalkotos
    class Partition(object):
45 de34e7aa Nikos Skalkotos
        """Represents a partition entry in MBR"""
46 de34e7aa Nikos Skalkotos
        format = "<B3sB3sLL"
47 de34e7aa Nikos Skalkotos
48 de34e7aa Nikos Skalkotos
        def __init__(self, raw_part):
49 de34e7aa Nikos Skalkotos
            """Create a Partition instance"""
50 de34e7aa Nikos Skalkotos
            (
51 de34e7aa Nikos Skalkotos
                self.status,
52 de34e7aa Nikos Skalkotos
                self.start,
53 de34e7aa Nikos Skalkotos
                self.type,
54 de34e7aa Nikos Skalkotos
                self.end,
55 de34e7aa Nikos Skalkotos
                self.first_sector,
56 de34e7aa Nikos Skalkotos
                self.sector_count
57 de34e7aa Nikos Skalkotos
            ) = struct.unpack(self.format, raw_part)
58 de34e7aa Nikos Skalkotos
59 de34e7aa Nikos Skalkotos
        def pack(self):
60 de34e7aa Nikos Skalkotos
            """Pack the partition values into a binary string"""
61 de34e7aa Nikos Skalkotos
            return struct.pack(self.format,
62 de34e7aa Nikos Skalkotos
                               self.status,
63 de34e7aa Nikos Skalkotos
                               self.start,
64 de34e7aa Nikos Skalkotos
                               self.type,
65 de34e7aa Nikos Skalkotos
                               self.end,
66 de34e7aa Nikos Skalkotos
                               self.first_sector,
67 de34e7aa Nikos Skalkotos
                               self.sector_count)
68 de34e7aa Nikos Skalkotos
69 de34e7aa Nikos Skalkotos
        @staticmethod
70 de34e7aa Nikos Skalkotos
        def size():
71 de34e7aa Nikos Skalkotos
            """Returns the size of an MBR partition entry"""
72 de34e7aa Nikos Skalkotos
            return struct.calcsize(MBR.Partition.format)
73 de34e7aa Nikos Skalkotos
74 de34e7aa Nikos Skalkotos
        def __str__(self):
75 de34e7aa Nikos Skalkotos
            start = self.unpack_chs(self.start)
76 de34e7aa Nikos Skalkotos
            end = self.unpack_chs(self.end)
77 de34e7aa Nikos Skalkotos
            return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
78 de34e7aa Nikos Skalkotos
                                          self.first_sector, self.sector_count)
79 de34e7aa Nikos Skalkotos
80 de34e7aa Nikos Skalkotos
        @staticmethod
81 de34e7aa Nikos Skalkotos
        def unpack_chs(chs):
82 de34e7aa Nikos Skalkotos
            """Unpacks a CHS address string to a tuple."""
83 de34e7aa Nikos Skalkotos
84 de34e7aa Nikos Skalkotos
            assert len(chs) == 3
85 de34e7aa Nikos Skalkotos
86 de34e7aa Nikos Skalkotos
            head = struct.unpack('<B', chs[0])[0]
87 de34e7aa Nikos Skalkotos
            sector = struct.unpack('<B', chs[1])[0] & 0x3f
88 de34e7aa Nikos Skalkotos
            cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
89 de34e7aa Nikos Skalkotos
                struct.unpack('<B', chs[2])[0]
90 de34e7aa Nikos Skalkotos
91 de34e7aa Nikos Skalkotos
            return (cylinder, head, sector)
92 de34e7aa Nikos Skalkotos
93 de34e7aa Nikos Skalkotos
        @staticmethod
94 de34e7aa Nikos Skalkotos
        def pack_chs(cylinder, head, sector):
95 de34e7aa Nikos Skalkotos
            """Packs a CHS tuple to an address string."""
96 de34e7aa Nikos Skalkotos
97 de34e7aa Nikos Skalkotos
            assert 1 <= sector <= 63
98 de34e7aa Nikos Skalkotos
            assert 0 <= head <= 255
99 0d413fc6 Nikos Skalkotos
            assert 0 <= cylinder
100 0d413fc6 Nikos Skalkotos
101 0d413fc6 Nikos Skalkotos
            # If the cylinders overflow then put the value (1023, 254, 63) to
102 0d413fc6 Nikos Skalkotos
            # the tuple. At least this is what OpenBSD does.
103 0d413fc6 Nikos Skalkotos
            if cylinder > 1023:
104 0d413fc6 Nikos Skalkotos
                cylinder = 1023
105 0d413fc6 Nikos Skalkotos
                head = 254
106 0d413fc6 Nikos Skalkotos
                sector = 63
107 de34e7aa Nikos Skalkotos
108 de34e7aa Nikos Skalkotos
            byte0 = head
109 de34e7aa Nikos Skalkotos
            byte1 = (cylinder >> 2) & 0xC0 | sector
110 de34e7aa Nikos Skalkotos
            byte2 = cylinder & 0xff
111 de34e7aa Nikos Skalkotos
112 de34e7aa Nikos Skalkotos
            return struct.pack('<BBB', byte0, byte1, byte2)
113 de34e7aa Nikos Skalkotos
114 de34e7aa Nikos Skalkotos
    format = "<444s2x16s16s16s16s2s"
115 de34e7aa Nikos Skalkotos
    """
116 de34e7aa Nikos Skalkotos
    Offset  Length          Contents
117 de34e7aa Nikos Skalkotos
    0       440(max. 446)   code area
118 de34e7aa Nikos Skalkotos
    440     2(optional)     disk signature
119 de34e7aa Nikos Skalkotos
    444     2               Usually nulls
120 de34e7aa Nikos Skalkotos
    446     16              Partition 0
121 de34e7aa Nikos Skalkotos
    462     16              Partition 1
122 de34e7aa Nikos Skalkotos
    478     16              Partition 2
123 de34e7aa Nikos Skalkotos
    494     16              Partition 3
124 de34e7aa Nikos Skalkotos
    510     2               MBR signature
125 de34e7aa Nikos Skalkotos
    """
126 de34e7aa Nikos Skalkotos
    def __init__(self, block):
127 de34e7aa Nikos Skalkotos
        """Create an MBR instance"""
128 de34e7aa Nikos Skalkotos
        raw_part = {}
129 de34e7aa Nikos Skalkotos
        (self.code_area,
130 de34e7aa Nikos Skalkotos
         raw_part[0],
131 de34e7aa Nikos Skalkotos
         raw_part[1],
132 de34e7aa Nikos Skalkotos
         raw_part[2],
133 de34e7aa Nikos Skalkotos
         raw_part[3],
134 de34e7aa Nikos Skalkotos
         self.signature) = struct.unpack(self.format, block)
135 de34e7aa Nikos Skalkotos
136 de34e7aa Nikos Skalkotos
        self.part = {}
137 de34e7aa Nikos Skalkotos
        for i in range(4):
138 de34e7aa Nikos Skalkotos
            self.part[i] = self.Partition(raw_part[i])
139 de34e7aa Nikos Skalkotos
140 de34e7aa Nikos Skalkotos
    @staticmethod
141 de34e7aa Nikos Skalkotos
    def size():
142 de34e7aa Nikos Skalkotos
        """Return the size of a Master Boot Record."""
143 de34e7aa Nikos Skalkotos
        return struct.calcsize(MBR.format)
144 de34e7aa Nikos Skalkotos
145 de34e7aa Nikos Skalkotos
    def pack(self):
146 de34e7aa Nikos Skalkotos
        """Pack an MBR to a binary string."""
147 de34e7aa Nikos Skalkotos
        return struct.pack(self.format,
148 de34e7aa Nikos Skalkotos
                           self.code_area,
149 de34e7aa Nikos Skalkotos
                           self.part[0].pack(),
150 de34e7aa Nikos Skalkotos
                           self.part[1].pack(),
151 de34e7aa Nikos Skalkotos
                           self.part[2].pack(),
152 de34e7aa Nikos Skalkotos
                           self.part[3].pack(),
153 de34e7aa Nikos Skalkotos
                           self.signature)
154 de34e7aa Nikos Skalkotos
155 de34e7aa Nikos Skalkotos
    def __str__(self):
156 de34e7aa Nikos Skalkotos
        ret = ""
157 de34e7aa Nikos Skalkotos
        for i in range(4):
158 de34e7aa Nikos Skalkotos
            ret += "Partition %d: %s\n" % (i, self.part[i])
159 de34e7aa Nikos Skalkotos
        ret += "Signature: %s %s\n" % (hex(ord(self.signature[0])),
160 de34e7aa Nikos Skalkotos
                                       hex(ord(self.signature[1])))
161 de34e7aa Nikos Skalkotos
        return ret
162 de34e7aa Nikos Skalkotos
163 de34e7aa Nikos Skalkotos
164 de34e7aa Nikos Skalkotos
class Disklabel:
165 de34e7aa Nikos Skalkotos
    """Represents an OpenBSD Disklabel"""
166 de34e7aa Nikos Skalkotos
    format = "<IHH16s16sIIIIII8sIHHIII20sHH16sIHHII364s"
167 de34e7aa Nikos Skalkotos
    """
168 de34e7aa Nikos Skalkotos
    Offset  Length          Contents
169 de34e7aa Nikos Skalkotos
    0       4               Magic
170 de34e7aa Nikos Skalkotos
    4       2               Drive Type
171 de34e7aa Nikos Skalkotos
    6       2               Subtype
172 de34e7aa Nikos Skalkotos
    8       16              Type Name
173 de34e7aa Nikos Skalkotos
    24      16              Pack Identifier
174 de34e7aa Nikos Skalkotos
    32      4               Bytes per sector
175 de34e7aa Nikos Skalkotos
    36      4               Data sectors per track
176 de34e7aa Nikos Skalkotos
    40      4               Tracks per cilinder
177 de34e7aa Nikos Skalkotos
    44      4               Data cylinders per unit
178 de34e7aa Nikos Skalkotos
    48      4               Data sectors per cylynder
179 de34e7aa Nikos Skalkotos
    52      4               Data sectors per unit
180 de34e7aa Nikos Skalkotos
    56      8               Unique label identifier
181 de34e7aa Nikos Skalkotos
    64      4               Alt cylinders per unit
182 de34e7aa Nikos Skalkotos
    68      2               Start of useable region (high part)
183 de34e7aa Nikos Skalkotos
    70      2               Size of usable region (high part)
184 de34e7aa Nikos Skalkotos
    72      4               Start of useable region
185 de34e7aa Nikos Skalkotos
    76      4               End of usable region
186 de34e7aa Nikos Skalkotos
    80      4               Generic Flags
187 de34e7aa Nikos Skalkotos
    84      5*4             Drive-type specific information
188 de34e7aa Nikos Skalkotos
    104     2               Number of data sectors (high part)
189 de34e7aa Nikos Skalkotos
    106     2               Version
190 de34e7aa Nikos Skalkotos
    108     4*4             Reserved for future use
191 de34e7aa Nikos Skalkotos
    124     4               Magic number
192 de34e7aa Nikos Skalkotos
    128     2               Xor of data Inclu. partitions
193 de34e7aa Nikos Skalkotos
    130     2               Number of partitions in following
194 de34e7aa Nikos Skalkotos
    132     4               size of boot area at sn0, bytes
195 de34e7aa Nikos Skalkotos
    136     4               Max size of fs superblock, bytes
196 de34e7aa Nikos Skalkotos
    140     16*16           Partition Table
197 de34e7aa Nikos Skalkotos
    """
198 de34e7aa Nikos Skalkotos
199 de34e7aa Nikos Skalkotos
    class PartitionTable:
200 de34e7aa Nikos Skalkotos
        """Reprepsents an OpenBSD Partition Table"""
201 de34e7aa Nikos Skalkotos
        format = "<IIHHBBH"
202 de34e7aa Nikos Skalkotos
        """
203 de34e7aa Nikos Skalkotos
        Partition Entry:
204 de34e7aa Nikos Skalkotos
        Offset  Length          Contents
205 de34e7aa Nikos Skalkotos
        0       4               Number of sectors in the partition
206 de34e7aa Nikos Skalkotos
        4       4               Starting sector
207 de34e7aa Nikos Skalkotos
        8       2               Starting sector (high part)
208 de34e7aa Nikos Skalkotos
        10      2               Number of sectors (high part)
209 de34e7aa Nikos Skalkotos
        12      1               Filesystem type
210 de34e7aa Nikos Skalkotos
        13      1               Filesystem Fragment per block
211 de34e7aa Nikos Skalkotos
        14      2               FS cylinders per group
212 de34e7aa Nikos Skalkotos
        """
213 de34e7aa Nikos Skalkotos
214 de34e7aa Nikos Skalkotos
        Partition = namedtuple(
215 de34e7aa Nikos Skalkotos
            'Partition', 'size, offset, offseth, sizeh, fstype, frag, cpg')
216 de34e7aa Nikos Skalkotos
217 de34e7aa Nikos Skalkotos
        def __init__(self, ptable, pnumber):
218 de34e7aa Nikos Skalkotos
            """Create a Partition Table instance"""
219 de34e7aa Nikos Skalkotos
            self.part = []
220 de34e7aa Nikos Skalkotos
221 de34e7aa Nikos Skalkotos
            size = struct.calcsize(self.format)
222 de34e7aa Nikos Skalkotos
223 de34e7aa Nikos Skalkotos
            raw = cStringIO.StringIO(ptable)
224 de34e7aa Nikos Skalkotos
            try:
225 de34e7aa Nikos Skalkotos
                for i in range(pnumber):
226 de34e7aa Nikos Skalkotos
                    p = self.Partition(
227 de34e7aa Nikos Skalkotos
                        *struct.unpack(self.format, raw.read(size)))
228 de34e7aa Nikos Skalkotos
                    self.part.append(p)
229 de34e7aa Nikos Skalkotos
            finally:
230 de34e7aa Nikos Skalkotos
                raw.close()
231 de34e7aa Nikos Skalkotos
232 de34e7aa Nikos Skalkotos
        def __str__(self):
233 de34e7aa Nikos Skalkotos
            """Print the Partition table"""
234 de34e7aa Nikos Skalkotos
            val = ""
235 de34e7aa Nikos Skalkotos
            for i in range(len(self.part)):
236 1bda0902 Nikos Skalkotos
                val += "%c: %s\n" % (chr(ord('a') + i), str(self.part[i]))
237 de34e7aa Nikos Skalkotos
            return val
238 de34e7aa Nikos Skalkotos
239 de34e7aa Nikos Skalkotos
        def pack(self):
240 de34e7aa Nikos Skalkotos
            """Packs the partition table into a binary string."""
241 de34e7aa Nikos Skalkotos
            ret = ""
242 de34e7aa Nikos Skalkotos
            for i in range(len(self.part)):
243 de34e7aa Nikos Skalkotos
                ret += struct.pack(self.format,
244 de34e7aa Nikos Skalkotos
                                   self.part[i].size,
245 de34e7aa Nikos Skalkotos
                                   self.part[i].offset,
246 de34e7aa Nikos Skalkotos
                                   self.part[i].offseth,
247 de34e7aa Nikos Skalkotos
                                   self.part[i].sizeh,
248 de34e7aa Nikos Skalkotos
                                   self.part[i].fstype,
249 de34e7aa Nikos Skalkotos
                                   self.part[i].frag,
250 de34e7aa Nikos Skalkotos
                                   self.part[i].cpg)
251 de34e7aa Nikos Skalkotos
            return ret
252 de34e7aa Nikos Skalkotos
253 de34e7aa Nikos Skalkotos
        def setpsize(self, i, size):
254 de34e7aa Nikos Skalkotos
            """Set size for partition i"""
255 de34e7aa Nikos Skalkotos
            tmp = self.part[i]
256 de34e7aa Nikos Skalkotos
            self.part[i] = self.Partition(size & 0xffffffff, tmp.offset,
257 de34e7aa Nikos Skalkotos
                                          tmp.offseth, size >> 32, tmp.fstype,
258 de34e7aa Nikos Skalkotos
                                          tmp.frag, tmp.cpg)
259 de34e7aa Nikos Skalkotos
260 de34e7aa Nikos Skalkotos
        def getpsize(self, i):
261 de34e7aa Nikos Skalkotos
            return (self.part[i].sizeh << 32) + self.part[i].size
262 de34e7aa Nikos Skalkotos
263 de34e7aa Nikos Skalkotos
        def setpoffset(self, i, offset):
264 de34e7aa Nikos Skalkotos
            """Set  offset for partition i"""
265 de34e7aa Nikos Skalkotos
            tmp = self.part[i]
266 de34e7aa Nikos Skalkotos
            self.part[i] = self.Partition(tmp.size, offset & 0xffffffff,
267 de34e7aa Nikos Skalkotos
                                          offset >> 32, tmp.sizeh, tmp.frag,
268 de34e7aa Nikos Skalkotos
                                          tmp.cpg)
269 de34e7aa Nikos Skalkotos
270 de34e7aa Nikos Skalkotos
        def getpoffset(self, i):
271 de34e7aa Nikos Skalkotos
            return (self.part[i].offseth << 32) + self.part[i].offset
272 de34e7aa Nikos Skalkotos
273 de34e7aa Nikos Skalkotos
    def __init__(self, disk):
274 de34e7aa Nikos Skalkotos
        """Create a DiskLabel instance"""
275 de34e7aa Nikos Skalkotos
        self.disk = disk
276 de34e7aa Nikos Skalkotos
        self.part_num = None
277 de34e7aa Nikos Skalkotos
278 de34e7aa Nikos Skalkotos
        with open(disk, "rb") as d:
279 de34e7aa Nikos Skalkotos
            sector0 = d.read(BLOCKSIZE)
280 de34e7aa Nikos Skalkotos
            self.mbr = MBR(sector0)
281 de34e7aa Nikos Skalkotos
282 de34e7aa Nikos Skalkotos
            for i in range(4):
283 de34e7aa Nikos Skalkotos
                if self.mbr.part[i].type == 0xa6:  # OpenBSD type
284 de34e7aa Nikos Skalkotos
                    self.part_num = i
285 de34e7aa Nikos Skalkotos
                    break
286 de34e7aa Nikos Skalkotos
287 de34e7aa Nikos Skalkotos
            assert self.part_num is not None, "No OpenBSD partition found"
288 de34e7aa Nikos Skalkotos
289 de34e7aa Nikos Skalkotos
            d.seek(BLOCKSIZE * self.mbr.part[self.part_num].first_sector)
290 de34e7aa Nikos Skalkotos
            part_sector0 = d.read(BLOCKSIZE)
291 de34e7aa Nikos Skalkotos
            # The offset of the disklabel from the begining of the
292 de34e7aa Nikos Skalkotos
            # partition is one sector
293 de34e7aa Nikos Skalkotos
            part_sector1 = d.read(BLOCKSIZE)
294 de34e7aa Nikos Skalkotos
295 de34e7aa Nikos Skalkotos
        (self.magic,
296 de34e7aa Nikos Skalkotos
         self.dtype,
297 de34e7aa Nikos Skalkotos
         self.subtype,
298 de34e7aa Nikos Skalkotos
         self.typename,
299 de34e7aa Nikos Skalkotos
         self.packname,
300 de34e7aa Nikos Skalkotos
         self.secsize,
301 de34e7aa Nikos Skalkotos
         self.nsectors,
302 de34e7aa Nikos Skalkotos
         self.ntracks,
303 de34e7aa Nikos Skalkotos
         self.ncylinders,
304 de34e7aa Nikos Skalkotos
         self.secpercyl,
305 de34e7aa Nikos Skalkotos
         self.secperunit,
306 de34e7aa Nikos Skalkotos
         self.uid,
307 de34e7aa Nikos Skalkotos
         self.acylinders,
308 de34e7aa Nikos Skalkotos
         self.bstarth,
309 de34e7aa Nikos Skalkotos
         self.bendh,
310 de34e7aa Nikos Skalkotos
         self.bstart,
311 de34e7aa Nikos Skalkotos
         self.bend,
312 de34e7aa Nikos Skalkotos
         self.flags,
313 de34e7aa Nikos Skalkotos
         self.drivedata,
314 de34e7aa Nikos Skalkotos
         self.secperunith,
315 de34e7aa Nikos Skalkotos
         self.version,
316 de34e7aa Nikos Skalkotos
         self.spare,
317 de34e7aa Nikos Skalkotos
         self.magic2,
318 de34e7aa Nikos Skalkotos
         self.checksum,
319 de34e7aa Nikos Skalkotos
         self.npartitions,
320 de34e7aa Nikos Skalkotos
         self.bbsize,
321 de34e7aa Nikos Skalkotos
         self.sbsize,
322 de34e7aa Nikos Skalkotos
         ptable_raw) = struct.unpack(self.format, part_sector1)
323 de34e7aa Nikos Skalkotos
324 de34e7aa Nikos Skalkotos
        assert self.magic == DISKMAGIC, "Disklabel is not valid"
325 de34e7aa Nikos Skalkotos
326 de34e7aa Nikos Skalkotos
        self.ptable = self.PartitionTable(ptable_raw, self.npartitions)
327 de34e7aa Nikos Skalkotos
328 de34e7aa Nikos Skalkotos
    def pack(self, checksum=None):
329 de34e7aa Nikos Skalkotos
        return struct.pack(self.format,
330 de34e7aa Nikos Skalkotos
                           self.magic,
331 de34e7aa Nikos Skalkotos
                           self.dtype,
332 de34e7aa Nikos Skalkotos
                           self.subtype,
333 de34e7aa Nikos Skalkotos
                           self.typename,
334 de34e7aa Nikos Skalkotos
                           self.packname,
335 de34e7aa Nikos Skalkotos
                           self.secsize,
336 de34e7aa Nikos Skalkotos
                           self.nsectors,
337 de34e7aa Nikos Skalkotos
                           self.ntracks,
338 de34e7aa Nikos Skalkotos
                           self.ncylinders,
339 de34e7aa Nikos Skalkotos
                           self.secpercyl,
340 de34e7aa Nikos Skalkotos
                           self.secperunit,
341 de34e7aa Nikos Skalkotos
                           self.uid,
342 de34e7aa Nikos Skalkotos
                           self.acylinders,
343 de34e7aa Nikos Skalkotos
                           self.bstarth,
344 de34e7aa Nikos Skalkotos
                           self.bendh,
345 de34e7aa Nikos Skalkotos
                           self.bstart,
346 de34e7aa Nikos Skalkotos
                           self.bend,
347 de34e7aa Nikos Skalkotos
                           self.flags,
348 de34e7aa Nikos Skalkotos
                           self.drivedata,
349 de34e7aa Nikos Skalkotos
                           self.secperunith,
350 de34e7aa Nikos Skalkotos
                           self.version,
351 de34e7aa Nikos Skalkotos
                           self.spare,
352 de34e7aa Nikos Skalkotos
                           self.magic2,
353 de34e7aa Nikos Skalkotos
                           self.checksum if checksum is None else checksum,
354 de34e7aa Nikos Skalkotos
                           self.npartitions,
355 de34e7aa Nikos Skalkotos
                           self.bbsize,
356 de34e7aa Nikos Skalkotos
                           self.sbsize,
357 de34e7aa Nikos Skalkotos
                           self.ptable.pack() +
358 ee78390c Nikos Skalkotos
                           ((364 - self.npartitions * 16) * '\x00'))
359 de34e7aa Nikos Skalkotos
360 de34e7aa Nikos Skalkotos
    def compute_checksum(self):
361 de34e7aa Nikos Skalkotos
        """Compute the checksum of the disklabel"""
362 de34e7aa Nikos Skalkotos
363 de34e7aa Nikos Skalkotos
        raw = cStringIO.StringIO(self.pack(0))
364 de34e7aa Nikos Skalkotos
        checksum = 0
365 de34e7aa Nikos Skalkotos
        try:
366 de34e7aa Nikos Skalkotos
            uint16 = raw.read(2)
367 de34e7aa Nikos Skalkotos
            while uint16 != "":
368 de34e7aa Nikos Skalkotos
                checksum ^= struct.unpack('<H', uint16)[0]
369 de34e7aa Nikos Skalkotos
                uint16 = raw.read(2)
370 de34e7aa Nikos Skalkotos
        finally:
371 de34e7aa Nikos Skalkotos
            raw.close()
372 de34e7aa Nikos Skalkotos
373 de34e7aa Nikos Skalkotos
        return checksum
374 de34e7aa Nikos Skalkotos
375 de34e7aa Nikos Skalkotos
    def setdsize(self, dsize):
376 de34e7aa Nikos Skalkotos
        """Set disk size"""
377 de34e7aa Nikos Skalkotos
        self.secperunith = dsize >> 32
378 de34e7aa Nikos Skalkotos
        self.secperunit = dsize & 0xffffffff
379 de34e7aa Nikos Skalkotos
380 de34e7aa Nikos Skalkotos
    def getdsize(self):
381 de34e7aa Nikos Skalkotos
        """Get disk size"""
382 de34e7aa Nikos Skalkotos
        return (self.secperunith << 32) + self.secperunit
383 de34e7aa Nikos Skalkotos
384 de34e7aa Nikos Skalkotos
    def setbstart(self, bstart):
385 de34e7aa Nikos Skalkotos
        """Set start of useable region"""
386 de34e7aa Nikos Skalkotos
        self.bstarth = bstart >> 32
387 de34e7aa Nikos Skalkotos
        self.bstart = bstart & 0xffffffff
388 de34e7aa Nikos Skalkotos
389 de34e7aa Nikos Skalkotos
    def getbstart(self):
390 de34e7aa Nikos Skalkotos
        """Get start of usable region"""
391 de34e7aa Nikos Skalkotos
        return (self.bstarth << 32) + self.bstart
392 de34e7aa Nikos Skalkotos
393 de34e7aa Nikos Skalkotos
    def setbend(self, bend):
394 de34e7aa Nikos Skalkotos
        """Set end of useable region"""
395 de34e7aa Nikos Skalkotos
        self.bendh = bend >> 32
396 de34e7aa Nikos Skalkotos
        self.bend = bend & 0xffffffff
397 de34e7aa Nikos Skalkotos
398 de34e7aa Nikos Skalkotos
    def getbend(self):
399 de34e7aa Nikos Skalkotos
        return (self.bendh << 32) + self.bend
400 de34e7aa Nikos Skalkotos
401 de34e7aa Nikos Skalkotos
    def enlarge_disk(self, new_size):
402 de34e7aa Nikos Skalkotos
        """Enlarge the size of the disk"""
403 de34e7aa Nikos Skalkotos
404 de34e7aa Nikos Skalkotos
        assert new_size >= self.secperunit, \
405 de34e7aa Nikos Skalkotos
            "New size cannot be smaller that %s" % self.secperunit
406 de34e7aa Nikos Skalkotos
407 de34e7aa Nikos Skalkotos
        # Fix the disklabel
408 de34e7aa Nikos Skalkotos
        self.setdsize(new_size)
409 de34e7aa Nikos Skalkotos
        self.ncylinders = self.getdsize() // (self.nsectors * self.ntracks)
410 de34e7aa Nikos Skalkotos
        self.setbend(self.ncylinders * self.nsectors * self.ntracks)
411 de34e7aa Nikos Skalkotos
412 de34e7aa Nikos Skalkotos
        # Partition 'c' descriptes the entire disk
413 de34e7aa Nikos Skalkotos
        self.ptable.setpsize(2, new_size)
414 de34e7aa Nikos Skalkotos
415 de34e7aa Nikos Skalkotos
        # Fix the MBR table
416 de34e7aa Nikos Skalkotos
        start = self.mbr.part[self.part_num].first_sector
417 de34e7aa Nikos Skalkotos
        self.mbr.part[self.part_num].sector_count = self.getbend() - start
418 de34e7aa Nikos Skalkotos
419 de34e7aa Nikos Skalkotos
        lba = self.getbend() - 1
420 de34e7aa Nikos Skalkotos
        cylinder = lba // (self.ntracks * self.nsectors)
421 de34e7aa Nikos Skalkotos
        header = (lba // self.nsectors) % self.ntracks
422 de34e7aa Nikos Skalkotos
        sector = (lba % self.nsectors) + 1
423 de34e7aa Nikos Skalkotos
        chs = MBR.Partition.pack_chs(cylinder, header, sector)
424 de34e7aa Nikos Skalkotos
        self.mbr.part[self.part_num].end = chs
425 de34e7aa Nikos Skalkotos
426 de34e7aa Nikos Skalkotos
        self.checksum = self.compute_checksum()
427 de34e7aa Nikos Skalkotos
428 de34e7aa Nikos Skalkotos
    def write(self):
429 de34e7aa Nikos Skalkotos
        """Write the disklabel back to the media"""
430 de34e7aa Nikos Skalkotos
        with open(self.disk, 'rw+b') as d:
431 de34e7aa Nikos Skalkotos
            d.write(self.mbr.pack())
432 de34e7aa Nikos Skalkotos
433 de34e7aa Nikos Skalkotos
            d.seek((self.mbr.part[self.part_num].first_sector + 1) * BLOCKSIZE)
434 de34e7aa Nikos Skalkotos
            d.write(self.pack())
435 de34e7aa Nikos Skalkotos
436 de34e7aa Nikos Skalkotos
    def get_last_partition_id(self):
437 de34e7aa Nikos Skalkotos
        """Returns the id of the last partition"""
438 de34e7aa Nikos Skalkotos
        part = 0
439 de34e7aa Nikos Skalkotos
        end = 0
440 de34e7aa Nikos Skalkotos
        # Don't check partition 'c' which is the whole disk
441 de34e7aa Nikos Skalkotos
        for i in filter(lambda x: x != 2, range(self.npartitions)):
442 de34e7aa Nikos Skalkotos
            curr_end = self.ptable.getpsize(i) + self.ptable.getpoffset(i)
443 de34e7aa Nikos Skalkotos
            if end < curr_end:
444 de34e7aa Nikos Skalkotos
                end = curr_end
445 de34e7aa Nikos Skalkotos
                part = i
446 de34e7aa Nikos Skalkotos
447 de34e7aa Nikos Skalkotos
        assert end > 0, "No partition found"
448 de34e7aa Nikos Skalkotos
449 de34e7aa Nikos Skalkotos
        return part
450 de34e7aa Nikos Skalkotos
451 de34e7aa Nikos Skalkotos
    def enlarge_last_partition(self):
452 de34e7aa Nikos Skalkotos
        """Enlarge the last partition to cover up all the free space"""
453 de34e7aa Nikos Skalkotos
454 de34e7aa Nikos Skalkotos
        part_num = self.get_last_partition_id()
455 de34e7aa Nikos Skalkotos
456 de34e7aa Nikos Skalkotos
        end = self.ptable.getpsize(part_num) + self.ptable.getpoffset(part_num)
457 de34e7aa Nikos Skalkotos
458 de34e7aa Nikos Skalkotos
        assert end > 0, "No partition found"
459 de34e7aa Nikos Skalkotos
460 de34e7aa Nikos Skalkotos
        if self.ptable.part[part_num].fstype == 1:  # Swap partition.
461 de34e7aa Nikos Skalkotos
            #TODO: Maybe create a warning?
462 de34e7aa Nikos Skalkotos
            return
463 de34e7aa Nikos Skalkotos
464 de34e7aa Nikos Skalkotos
        if end > (self.getbend() - 1024):
465 de34e7aa Nikos Skalkotos
            return
466 de34e7aa Nikos Skalkotos
467 de34e7aa Nikos Skalkotos
        self.ptable.setpsize(
468 de34e7aa Nikos Skalkotos
            part_num, self.getbend() - self.ptable.getpoffset(part_num) - 1024)
469 de34e7aa Nikos Skalkotos
470 de34e7aa Nikos Skalkotos
        self.checksum = self.compute_checksum()
471 de34e7aa Nikos Skalkotos
472 de34e7aa Nikos Skalkotos
    def __str__(self):
473 de34e7aa Nikos Skalkotos
        """Print the Disklabel"""
474 de34e7aa Nikos Skalkotos
        title1 = "Master Boot Record"
475 de34e7aa Nikos Skalkotos
        title2 = "Disklabel"
476 de34e7aa Nikos Skalkotos
477 de34e7aa Nikos Skalkotos
        return \
478 de34e7aa Nikos Skalkotos
            "%s\n%s\n%s\n" % (title1, len(title1) * "=", str(self.mbr)) + \
479 de34e7aa Nikos Skalkotos
            "%s\n%s\n" % (title2, len(title2) * "=") + \
480 de34e7aa Nikos Skalkotos
            "Magic Number: 0x%x\n" % self.magic + \
481 de34e7aa Nikos Skalkotos
            "Drive type: %d\n" % self.dtype + \
482 de34e7aa Nikos Skalkotos
            "Subtype: %d\n" % self.subtype + \
483 abe2770c Nikos Skalkotos
            "Typename: %s\n" % self.typename.strip('\x00').strip() + \
484 abe2770c Nikos Skalkotos
            "Pack Identifier: %s\n" % self.packname.strip('\x00').strip() + \
485 de34e7aa Nikos Skalkotos
            "Number of bytes per sector: %d\n" % self.secsize + \
486 de34e7aa Nikos Skalkotos
            "Number of data sectors per track: %d\n" % self.nsectors + \
487 de34e7aa Nikos Skalkotos
            "Number of tracks per cylinder: %d\n" % self.ntracks + \
488 de34e7aa Nikos Skalkotos
            "Number of data cylinders per unit: %d\n" % self.ncylinders + \
489 de34e7aa Nikos Skalkotos
            "Number of data sectors per cylinder: %d\n" % self.secpercyl + \
490 de34e7aa Nikos Skalkotos
            "Number of data sectors per unit: %d\n" % self.secperunit + \
491 de34e7aa Nikos Skalkotos
            "DUID: %s\n" % "".join(x.encode('hex') for x in self.uid) + \
492 de34e7aa Nikos Skalkotos
            "Alt. cylinders per unit: %d\n" % self.acylinders + \
493 de34e7aa Nikos Skalkotos
            "Start of useable region (high part): %d\n" % self.bstarth + \
494 de34e7aa Nikos Skalkotos
            "Size of useable region (high part): %d\n" % self.bendh + \
495 de34e7aa Nikos Skalkotos
            "Start of useable region: %d\n" % self.bstart + \
496 de34e7aa Nikos Skalkotos
            "End of usable region: %d\n" % self.bend + \
497 de34e7aa Nikos Skalkotos
            "Generic Flags: %r\n" % self.flags + \
498 de34e7aa Nikos Skalkotos
            "Drive data: %r\n" % self.drivedata + \
499 de34e7aa Nikos Skalkotos
            "Number of data sectors (high part): %d\n" % self.secperunith + \
500 de34e7aa Nikos Skalkotos
            "Version: %d\n" % self.version + \
501 de34e7aa Nikos Skalkotos
            "Reserved for future use: %r\n" % self.spare + \
502 de34e7aa Nikos Skalkotos
            "The magic number again: 0x%x\n" % self.magic2 + \
503 de34e7aa Nikos Skalkotos
            "Checksum: %d\n" % self.checksum + \
504 de34e7aa Nikos Skalkotos
            "Number of partitions: %d\n" % self.npartitions + \
505 de34e7aa Nikos Skalkotos
            "Size of boot aread at sn0: %d\n" % self.bbsize + \
506 de34e7aa Nikos Skalkotos
            "Max size of fs superblock: %d\n" % self.sbsize + \
507 de34e7aa Nikos Skalkotos
            "%s" % self.ptable
508 de34e7aa Nikos Skalkotos
509 de34e7aa Nikos Skalkotos
510 de34e7aa Nikos Skalkotos
if __name__ == '__main__':
511 de34e7aa Nikos Skalkotos
512 de34e7aa Nikos Skalkotos
    usage = "Usage: %prog [options] <input_media>"
513 de34e7aa Nikos Skalkotos
    parser = optparse.OptionParser(usage=usage)
514 de34e7aa Nikos Skalkotos
515 de34e7aa Nikos Skalkotos
    parser.add_option("-l", "--list", action="store_true", dest="list",
516 de34e7aa Nikos Skalkotos
                      default=False,
517 de34e7aa Nikos Skalkotos
                      help="list the disklabel on the specified media")
518 de34e7aa Nikos Skalkotos
    parser.add_option("--print-last", action="store_true", dest="last_part",
519 de34e7aa Nikos Skalkotos
                      default=False,
520 de34e7aa Nikos Skalkotos
                      help="print the label of the last partition")
521 de34e7aa Nikos Skalkotos
    parser.add_option("--print-last-linux", action="store_true",
522 de34e7aa Nikos Skalkotos
                      dest="last_linux", default=False,
523 de34e7aa Nikos Skalkotos
                      help="print the linux number for the last partition")
524 de34e7aa Nikos Skalkotos
    parser.add_option("--print-duid", action="store_true", dest="duid",
525 de34e7aa Nikos Skalkotos
                      default=False,
526 de34e7aa Nikos Skalkotos
                      help="print the disklabel unique identifier")
527 de34e7aa Nikos Skalkotos
    parser.add_option("-d", "--enlarge-disk", type="int", dest="disk_size",
528 de34e7aa Nikos Skalkotos
                      default=None, metavar="SIZE",
529 de34e7aa Nikos Skalkotos
                      help="Enlarge the disk to this SIZE (in sectors)")
530 de34e7aa Nikos Skalkotos
    parser.add_option(
531 de34e7aa Nikos Skalkotos
        "-p", "--enlarge-partition", action="store_true",
532 de34e7aa Nikos Skalkotos
        dest="enlarge_partition", default=False,
533 de34e7aa Nikos Skalkotos
        help="Enlarge the last partition to cover up the free space")
534 de34e7aa Nikos Skalkotos
535 de34e7aa Nikos Skalkotos
    options, args = parser.parse_args(sys.argv[1:])
536 de34e7aa Nikos Skalkotos
537 de34e7aa Nikos Skalkotos
    if len(args) != 1:
538 de34e7aa Nikos Skalkotos
        parser.error("Wrong number of arguments")
539 de34e7aa Nikos Skalkotos
540 de34e7aa Nikos Skalkotos
    disklabel = Disklabel(args[0])
541 de34e7aa Nikos Skalkotos
542 de34e7aa Nikos Skalkotos
    if options.list:
543 de34e7aa Nikos Skalkotos
        print disklabel
544 de34e7aa Nikos Skalkotos
        sys.exit(0)
545 de34e7aa Nikos Skalkotos
546 de34e7aa Nikos Skalkotos
    if options.duid:
547 de34e7aa Nikos Skalkotos
        print "%s" % "".join(x.encode('hex') for x in disklabel.uid)
548 de34e7aa Nikos Skalkotos
        sys.exit(0)
549 de34e7aa Nikos Skalkotos
550 de34e7aa Nikos Skalkotos
    if options.last_part:
551 de34e7aa Nikos Skalkotos
        print "%c" % chr(ord('a') + disklabel.get_last_partition_id())
552 de34e7aa Nikos Skalkotos
553 de34e7aa Nikos Skalkotos
    if options.last_linux:
554 de34e7aa Nikos Skalkotos
        part_id = disklabel.get_last_partition_id()
555 de34e7aa Nikos Skalkotos
        # The linux kernel does not assign a partition for label 'c' that
556 de34e7aa Nikos Skalkotos
        # describes the whole disk
557 de34e7aa Nikos Skalkotos
        print part_id + (4 if part_id > 2 else 5)
558 de34e7aa Nikos Skalkotos
559 de34e7aa Nikos Skalkotos
    if options.disk_size is not None:
560 de34e7aa Nikos Skalkotos
        disklabel.enlarge_disk(options.disk_size)
561 de34e7aa Nikos Skalkotos
562 de34e7aa Nikos Skalkotos
    if options.enlarge_partition:
563 de34e7aa Nikos Skalkotos
        disklabel.enlarge_last_partition()
564 de34e7aa Nikos Skalkotos
565 de34e7aa Nikos Skalkotos
    disklabel.write()
566 de34e7aa Nikos Skalkotos
567 de34e7aa Nikos Skalkotos
sys.exit(0)
568 de34e7aa Nikos Skalkotos
569 de34e7aa Nikos Skalkotos
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :