Statistics
| Branch: | Tag: | Revision:

root / snf-image-helper / fix_gpt.py @ 21be5a41

History | View | Annotate | Download (9.1 kB)

1 822ba661 Nikos Skalkotos
#!/usr/bin/env python
2 822ba661 Nikos Skalkotos
3 822ba661 Nikos Skalkotos
# Copyright (C) 2012 GRNET S.A.
4 822ba661 Nikos Skalkotos
#
5 822ba661 Nikos Skalkotos
# This program is free software; you can redistribute it and/or modify
6 822ba661 Nikos Skalkotos
# it under the terms of the GNU General Public License as published by
7 822ba661 Nikos Skalkotos
# the Free Software Foundation; either version 2 of the License, or
8 822ba661 Nikos Skalkotos
# (at your option) any later version.
9 822ba661 Nikos Skalkotos
#
10 822ba661 Nikos Skalkotos
# This program is distributed in the hope that it will be useful, but
11 822ba661 Nikos Skalkotos
# WITHOUT ANY WARRANTY; without even the implied warranty of
12 822ba661 Nikos Skalkotos
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
13 822ba661 Nikos Skalkotos
# General Public License for more details.
14 822ba661 Nikos Skalkotos
#
15 822ba661 Nikos Skalkotos
# You should have received a copy of the GNU General Public License
16 822ba661 Nikos Skalkotos
# along with this program; if not, write to the Free Software
17 822ba661 Nikos Skalkotos
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
18 822ba661 Nikos Skalkotos
# 02110-1301, USA.
19 822ba661 Nikos Skalkotos
20 822ba661 Nikos Skalkotos
import struct
21 822ba661 Nikos Skalkotos
import sys
22 822ba661 Nikos Skalkotos
import binascii
23 822ba661 Nikos Skalkotos
import stat
24 822ba661 Nikos Skalkotos
import os
25 822ba661 Nikos Skalkotos
26 822ba661 Nikos Skalkotos
BLOCKSIZE = 512
27 822ba661 Nikos Skalkotos
28 822ba661 Nikos Skalkotos
29 822ba661 Nikos Skalkotos
class MBR(object):
30 822ba661 Nikos Skalkotos
    """Represents a Master Boot Record."""
31 822ba661 Nikos Skalkotos
    class Partition(object):
32 822ba661 Nikos Skalkotos
        format = "<B3sB3sLL"
33 822ba661 Nikos Skalkotos
34 822ba661 Nikos Skalkotos
        def __init__(self, raw_part):
35 822ba661 Nikos Skalkotos
            (
36 822ba661 Nikos Skalkotos
                self.status,
37 822ba661 Nikos Skalkotos
                self.start,
38 822ba661 Nikos Skalkotos
                self.type,
39 822ba661 Nikos Skalkotos
                self.end,
40 822ba661 Nikos Skalkotos
                self.first_sector,
41 822ba661 Nikos Skalkotos
                self.sector_count
42 822ba661 Nikos Skalkotos
            ) = struct.unpack(self.format, raw_part)
43 822ba661 Nikos Skalkotos
44 822ba661 Nikos Skalkotos
        def pack(self):
45 822ba661 Nikos Skalkotos
            return struct.pack(self.format,
46 822ba661 Nikos Skalkotos
                self.status,
47 822ba661 Nikos Skalkotos
                self.start,
48 822ba661 Nikos Skalkotos
                self.type,
49 822ba661 Nikos Skalkotos
                self.end,
50 822ba661 Nikos Skalkotos
                self.first_sector,
51 822ba661 Nikos Skalkotos
                self.sector_count
52 822ba661 Nikos Skalkotos
            )
53 822ba661 Nikos Skalkotos
54 822ba661 Nikos Skalkotos
        @staticmethod
55 822ba661 Nikos Skalkotos
        def size():
56 822ba661 Nikos Skalkotos
            """Returns the size of an MBR partition entry"""
57 822ba661 Nikos Skalkotos
            return struct.calcsize(MBR.Partition.format)
58 822ba661 Nikos Skalkotos
59 822ba661 Nikos Skalkotos
        def unpack_chs(self, chs):
60 822ba661 Nikos Skalkotos
            """Unpacks a CHS address string to a tuple."""
61 822ba661 Nikos Skalkotos
62 822ba661 Nikos Skalkotos
            assert len(chs) == 3
63 822ba661 Nikos Skalkotos
64 822ba661 Nikos Skalkotos
            head = struct.unpack('<B', chs[0])[0]
65 822ba661 Nikos Skalkotos
            sector = struct.unpack('<B', chs[1])[0] & 0x3f
66 822ba661 Nikos Skalkotos
            cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
67 822ba661 Nikos Skalkotos
                struct.unpack('<B', chs[2])[0]
68 822ba661 Nikos Skalkotos
69 822ba661 Nikos Skalkotos
            return (cylinder, head, sector)
70 822ba661 Nikos Skalkotos
71 822ba661 Nikos Skalkotos
        def pack_chs(self, cylinder, head, sector):
72 822ba661 Nikos Skalkotos
            """Packs a CHS tuple to an address string."""
73 822ba661 Nikos Skalkotos
74 822ba661 Nikos Skalkotos
            assert 1 <= sector <= 63
75 822ba661 Nikos Skalkotos
            assert 0 <= cylinder <= 1023
76 822ba661 Nikos Skalkotos
            assert 0 <= head <= 255
77 822ba661 Nikos Skalkotos
78 822ba661 Nikos Skalkotos
            byte0 = head
79 822ba661 Nikos Skalkotos
            byte1 = (cylinder >> 2) & 0xC0 | sector
80 822ba661 Nikos Skalkotos
            byte2 = cylinder & 0xff
81 822ba661 Nikos Skalkotos
82 822ba661 Nikos Skalkotos
            return struct.pack('<BBB', byte0, byte1, byte2)
83 822ba661 Nikos Skalkotos
84 822ba661 Nikos Skalkotos
    format = "<444s2x16s16s16s16s2s"
85 822ba661 Nikos Skalkotos
    """
86 822ba661 Nikos Skalkotos
    Offset  Length          Contents
87 822ba661 Nikos Skalkotos
    0       440(max. 446)   code area
88 822ba661 Nikos Skalkotos
    440     2(optional)     disk signature
89 822ba661 Nikos Skalkotos
    444     2               Usually nulls
90 822ba661 Nikos Skalkotos
    446     16              Partition 0
91 822ba661 Nikos Skalkotos
    462     16              Partition 1
92 822ba661 Nikos Skalkotos
    478     16              Partition 2
93 822ba661 Nikos Skalkotos
    494     16              Partition 3
94 822ba661 Nikos Skalkotos
    510     2               MBR signature
95 822ba661 Nikos Skalkotos
    """
96 822ba661 Nikos Skalkotos
    def __init__(self, block):
97 822ba661 Nikos Skalkotos
        raw_part = {}
98 822ba661 Nikos Skalkotos
        self.code_area, \
99 822ba661 Nikos Skalkotos
        raw_part[0], \
100 822ba661 Nikos Skalkotos
        raw_part[1], \
101 822ba661 Nikos Skalkotos
        raw_part[2], \
102 822ba661 Nikos Skalkotos
        raw_part[3], \
103 822ba661 Nikos Skalkotos
        self.signature = struct.unpack(self.format, block)
104 822ba661 Nikos Skalkotos
105 822ba661 Nikos Skalkotos
        self.part = {}
106 822ba661 Nikos Skalkotos
        for i in range(4):
107 822ba661 Nikos Skalkotos
            self.part[i] = self.Partition(raw_part[i])
108 822ba661 Nikos Skalkotos
109 822ba661 Nikos Skalkotos
    @staticmethod
110 822ba661 Nikos Skalkotos
    def size():
111 822ba661 Nikos Skalkotos
        """Returns the size of a Master Boot Record."""
112 822ba661 Nikos Skalkotos
        return struct.calcsize(MBR.format)
113 822ba661 Nikos Skalkotos
114 822ba661 Nikos Skalkotos
    def pack(self):
115 822ba661 Nikos Skalkotos
        """Packs an MBR to a binary string."""
116 822ba661 Nikos Skalkotos
        return struct.pack(self.format,
117 822ba661 Nikos Skalkotos
            self.code_area,
118 822ba661 Nikos Skalkotos
            self.part[0].pack(),
119 822ba661 Nikos Skalkotos
            self.part[1].pack(),
120 822ba661 Nikos Skalkotos
            self.part[2].pack(),
121 822ba661 Nikos Skalkotos
            self.part[3].pack(),
122 822ba661 Nikos Skalkotos
            self.signature
123 822ba661 Nikos Skalkotos
        )
124 822ba661 Nikos Skalkotos
125 822ba661 Nikos Skalkotos
126 822ba661 Nikos Skalkotos
class GPTPartitionTable(object):
127 822ba661 Nikos Skalkotos
    """Represents a GUID Partition Table."""
128 822ba661 Nikos Skalkotos
    class GPTHeader(object):
129 822ba661 Nikos Skalkotos
        """Represents a GPT Header of a GUID Partition Table."""
130 822ba661 Nikos Skalkotos
        format = "<8s4sII4xQQQQ16sQIII"
131 822ba661 Nikos Skalkotos
        """
132 822ba661 Nikos Skalkotos
        Offset        Length                 Contents
133 822ba661 Nikos Skalkotos
        0       8 bytes         Signature
134 822ba661 Nikos Skalkotos
        8       4 bytes         Revision
135 822ba661 Nikos Skalkotos
        12      4 bytes         Header size in little endian
136 822ba661 Nikos Skalkotos
        16         4 bytes         CRC32 of header
137 822ba661 Nikos Skalkotos
        20         4 bytes         Reserved; must be zero
138 822ba661 Nikos Skalkotos
        24         8 bytes         Current LBA
139 822ba661 Nikos Skalkotos
        32         8 bytes         Backup LBA
140 822ba661 Nikos Skalkotos
        40         8 bytes         First usable LBA for partitions
141 822ba661 Nikos Skalkotos
        48         8 bytes         Last usable LBA
142 822ba661 Nikos Skalkotos
        56         16 bytes         Disk GUID
143 822ba661 Nikos Skalkotos
        72         8 bytes         Partition entries starting LBA
144 822ba661 Nikos Skalkotos
        80         4 bytes         Number of partition entries
145 822ba661 Nikos Skalkotos
        84         4 bytes         Size of a partition entry
146 822ba661 Nikos Skalkotos
        88         4 bytes         CRC32 of partition array
147 822ba661 Nikos Skalkotos
        92         *                 Reserved; must be zeroes
148 822ba661 Nikos Skalkotos
        LBA    size            Total
149 822ba661 Nikos Skalkotos
        """
150 822ba661 Nikos Skalkotos
151 822ba661 Nikos Skalkotos
        def __init__(self, block):
152 822ba661 Nikos Skalkotos
            self.signature, \
153 822ba661 Nikos Skalkotos
            self.revision, \
154 822ba661 Nikos Skalkotos
            self.hdr_size, \
155 822ba661 Nikos Skalkotos
            self.header_crc32, \
156 822ba661 Nikos Skalkotos
            self.current_lba, \
157 822ba661 Nikos Skalkotos
            self.backup_lba, \
158 822ba661 Nikos Skalkotos
            self.first_usable_lba, \
159 822ba661 Nikos Skalkotos
            self.last_usable_lba, \
160 822ba661 Nikos Skalkotos
            self.uuid, \
161 822ba661 Nikos Skalkotos
            self.part_entry_start, \
162 822ba661 Nikos Skalkotos
            self.part_count, \
163 822ba661 Nikos Skalkotos
            self.part_entry_size, \
164 822ba661 Nikos Skalkotos
            self.part_crc32 = struct.unpack(self.format, block)
165 822ba661 Nikos Skalkotos
166 822ba661 Nikos Skalkotos
        def pack(self):
167 822ba661 Nikos Skalkotos
            """Packs a GPT Header to a binary string."""
168 822ba661 Nikos Skalkotos
            return struct.pack(self.format,
169 822ba661 Nikos Skalkotos
                self.signature, \
170 822ba661 Nikos Skalkotos
                self.revision, \
171 822ba661 Nikos Skalkotos
                self.hdr_size, \
172 822ba661 Nikos Skalkotos
                self.header_crc32, \
173 822ba661 Nikos Skalkotos
                self.current_lba, \
174 822ba661 Nikos Skalkotos
                self.backup_lba, \
175 822ba661 Nikos Skalkotos
                self.first_usable_lba, \
176 822ba661 Nikos Skalkotos
                self.last_usable_lba, \
177 822ba661 Nikos Skalkotos
                self.uuid, \
178 822ba661 Nikos Skalkotos
                self.part_entry_start, \
179 822ba661 Nikos Skalkotos
                self.part_count, \
180 822ba661 Nikos Skalkotos
                self.part_entry_size, \
181 822ba661 Nikos Skalkotos
                self.part_crc32
182 822ba661 Nikos Skalkotos
            )
183 822ba661 Nikos Skalkotos
184 822ba661 Nikos Skalkotos
        @staticmethod
185 822ba661 Nikos Skalkotos
        def size():
186 822ba661 Nikos Skalkotos
            """Returns the size of a GPT Header."""
187 822ba661 Nikos Skalkotos
            return struct.calcsize(GPTPartitionTable.GPTHeader.format)
188 822ba661 Nikos Skalkotos
189 822ba661 Nikos Skalkotos
    def __init__(self, disk):
190 822ba661 Nikos Skalkotos
        self.disk = disk
191 822ba661 Nikos Skalkotos
        with open(disk, "rb") as d:
192 822ba661 Nikos Skalkotos
            # MBR (Logical block address 0)
193 822ba661 Nikos Skalkotos
            lba0 = d.read(BLOCKSIZE)
194 822ba661 Nikos Skalkotos
            self.mbr = MBR(lba0)
195 822ba661 Nikos Skalkotos
196 822ba661 Nikos Skalkotos
            # Primary GPT Header (LBA 1)
197 822ba661 Nikos Skalkotos
            raw_header = d.read(self.GPTHeader.size())
198 822ba661 Nikos Skalkotos
            self.primary = self.GPTHeader(raw_header)
199 822ba661 Nikos Skalkotos
200 822ba661 Nikos Skalkotos
            # Partition entries (LBA 2...34)
201 822ba661 Nikos Skalkotos
            d.seek(self.primary.part_entry_start * BLOCKSIZE)
202 822ba661 Nikos Skalkotos
            entries_size = self.primary.part_count * \
203 822ba661 Nikos Skalkotos
                                                self.primary.part_entry_size
204 822ba661 Nikos Skalkotos
            self.part_entries = d.read(entries_size)
205 822ba661 Nikos Skalkotos
206 822ba661 Nikos Skalkotos
            # Secondary GPT Header (LBA -1)
207 822ba661 Nikos Skalkotos
            d.seek(self.primary.backup_lba * BLOCKSIZE)
208 822ba661 Nikos Skalkotos
            raw_header = d.read(self.GPTHeader.size())
209 822ba661 Nikos Skalkotos
            self.secondary = self.GPTHeader(raw_header)
210 822ba661 Nikos Skalkotos
211 822ba661 Nikos Skalkotos
    def size(self):
212 822ba661 Nikos Skalkotos
        """Returns the payload size of GPT partitioned device."""
213 822ba661 Nikos Skalkotos
        return (self.primary.backup_lba + 1) * BLOCKSIZE
214 822ba661 Nikos Skalkotos
215 822ba661 Nikos Skalkotos
    def fix(self, lba_count):
216 822ba661 Nikos Skalkotos
        """Move the secondary GPT Header entries to the LBA specified by
217 822ba661 Nikos Skalkotos
        lba_count parameter.
218 822ba661 Nikos Skalkotos
        """
219 822ba661 Nikos Skalkotos
220 822ba661 Nikos Skalkotos
        assert lba_count * BLOCKSIZE > self.size()
221 822ba661 Nikos Skalkotos
222 822ba661 Nikos Skalkotos
        # Correct MBR
223 822ba661 Nikos Skalkotos
        #TODO: Check if the partition tables is hybrid
224 822ba661 Nikos Skalkotos
        self.mbr.part[0].sector_count = lba_count - 1
225 822ba661 Nikos Skalkotos
226 822ba661 Nikos Skalkotos
        # Fix Primary header
227 822ba661 Nikos Skalkotos
        self.primary.header_crc32 = 0
228 822ba661 Nikos Skalkotos
        self.primary.backup_lba = lba_count - 1  # LBA-1
229 822ba661 Nikos Skalkotos
        self.primary.last_usable_lba = lba_count - 34  # LBA-34
230 822ba661 Nikos Skalkotos
        self.primary.header_crc32 = \
231 822ba661 Nikos Skalkotos
                            binascii.crc32(self.primary.pack()) & 0xffffffff
232 822ba661 Nikos Skalkotos
233 822ba661 Nikos Skalkotos
        # Fix Secondary header
234 822ba661 Nikos Skalkotos
        self.secondary.header_crc32 = 0
235 822ba661 Nikos Skalkotos
        self.secondary.current_lba = self.primary.backup_lba
236 822ba661 Nikos Skalkotos
        self.secondary.last_usable_lba = lba_count - 34  # LBA-34
237 822ba661 Nikos Skalkotos
        self.secondary.part_entry_start = lba_count - 33  # LBA-33
238 822ba661 Nikos Skalkotos
        self.secondary.header_crc32 = \
239 822ba661 Nikos Skalkotos
                            binascii.crc32(self.secondary.pack()) & 0xffffffff
240 822ba661 Nikos Skalkotos
241 822ba661 Nikos Skalkotos
        # Copy the new partition table back to the device
242 822ba661 Nikos Skalkotos
        with open(self.disk, "wb") as d:
243 822ba661 Nikos Skalkotos
            d.write(self.mbr.pack())
244 822ba661 Nikos Skalkotos
            d.write(self.primary.pack())
245 822ba661 Nikos Skalkotos
            d.write('\x00' * (BLOCKSIZE - self.primary.size()))
246 822ba661 Nikos Skalkotos
            d.seek(self.secondary.part_entry_start * BLOCKSIZE)
247 822ba661 Nikos Skalkotos
            d.write(self.part_entries)
248 822ba661 Nikos Skalkotos
            d.seek(self.primary.backup_lba * BLOCKSIZE)
249 822ba661 Nikos Skalkotos
            d.write(self.secondary.pack())
250 822ba661 Nikos Skalkotos
            d.write('\x00' * (BLOCKSIZE - self.secondary.size()))
251 822ba661 Nikos Skalkotos
252 822ba661 Nikos Skalkotos
253 822ba661 Nikos Skalkotos
if __name__ == '__main__':
254 822ba661 Nikos Skalkotos
    usage = "Usage: %s <disk> <sectors>\n" % (sys.argv[0])
255 822ba661 Nikos Skalkotos
256 822ba661 Nikos Skalkotos
    if len(sys.argv) != 3:
257 822ba661 Nikos Skalkotos
        sys.stderr.write(usage)
258 822ba661 Nikos Skalkotos
        sys.exit(1)
259 822ba661 Nikos Skalkotos
260 822ba661 Nikos Skalkotos
    disk = sys.argv[1]
261 822ba661 Nikos Skalkotos
    mode = os.stat(disk).st_mode
262 822ba661 Nikos Skalkotos
    if not stat.S_ISBLK(mode):
263 822ba661 Nikos Skalkotos
        sys.stderr.write("Parameter disk must be a block device\n")
264 822ba661 Nikos Skalkotos
        sys.stderr.write(usage)
265 822ba661 Nikos Skalkotos
        sys.exit(1)
266 822ba661 Nikos Skalkotos
267 822ba661 Nikos Skalkotos
    try:
268 822ba661 Nikos Skalkotos
        size = int(sys.argv[2])
269 822ba661 Nikos Skalkotos
    except ValueError:
270 822ba661 Nikos Skalkotos
        sys.stderr.write("Parameter new_size must be a number\n")
271 822ba661 Nikos Skalkotos
        sys.stderr.write(usage)
272 822ba661 Nikos Skalkotos
        sys.exit(1)
273 822ba661 Nikos Skalkotos
274 822ba661 Nikos Skalkotos
    ptable = GPTPartitionTable(disk)
275 822ba661 Nikos Skalkotos
    if size * BLOCKSIZE == ptable.size():
276 822ba661 Nikos Skalkotos
        sys.stderr.write("Nothing to do...\n")
277 822ba661 Nikos Skalkotos
    elif size * BLOCKSIZE > ptable.size():
278 822ba661 Nikos Skalkotos
        ptable.fix(size)
279 822ba661 Nikos Skalkotos
        sys.stderr.write("GPT table was fixed\n")
280 822ba661 Nikos Skalkotos
    else:
281 822ba661 Nikos Skalkotos
        sys.stderr.write("Disk is langer than size")
282 822ba661 Nikos Skalkotos
        exit(1)
283 822ba661 Nikos Skalkotos
284 822ba661 Nikos Skalkotos
    sys.exit(0)
285 822ba661 Nikos Skalkotos
286 822ba661 Nikos Skalkotos
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :