~justin-fathomdb/nova/virtualbox-support

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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
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
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright 2010 United States Government as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""
Utility methods to resize, repartition, and modify disk images.
Includes injection of SSH PGP keys into authorized_keys file.
"""

import logging
import os
import tempfile

from nova import exception


def partition(infile, outfile, local_bytes=0, local_type='ext2', execute=None):
    """Takes a single partition represented by infile and writes a bootable
    drive image into outfile.

    The first 63 sectors (0-62) of the resulting image is a master boot record.
    Infile becomes the first primary partition.
    If local bytes is specified, a second primary partition is created and
    formatted as ext2.

    In the diagram below, dashes represent drive sectors.
    +-----+------. . .-------+------. . .------+
    | 0  a| b               c|d               e|
    +-----+------. . .-------+------. . .------+
    | mbr | primary partiton | local partition |
    +-----+------. . .-------+------. . .------+
    """
    sector_size = 512
    file_size = os.path.getsize(infile)
    if file_size % sector_size != 0:
        logging.warn("Input partition size not evenly divisible by"
                     " sector size: %d / %d", file_size, sector_size)
    primary_sectors = file_size / sector_size
    if local_bytes % sector_size != 0:
        logging.warn("Bytes for local storage not evenly divisible"
                     " by sector size: %d / %d", local_bytes, sector_size)
    local_sectors = local_bytes / sector_size

    mbr_last = 62 # a
    primary_first = mbr_last + 1 # b
    primary_last = primary_first + primary_sectors - 1 # c
    local_first = primary_last + 1 # d
    local_last = local_first + local_sectors - 1 # e
    last_sector = local_last # e

    # create an empty file
    execute('dd if=/dev/zero of=%s count=1 seek=%d bs=%d'
                  % (outfile, last_sector, sector_size))

    # make mbr partition
    execute('parted --script %s mklabel msdos' % outfile)

    # make primary partition
    execute('parted --script %s mkpart primary %ds %ds'
                  % (outfile, primary_first, primary_last))

    # make local partition
    if local_bytes > 0:
        execute('parted --script %s mkpartfs primary %s %ds %ds'
                      % (outfile, local_type, local_first, local_last))

    # copy file into partition
    execute('dd if=%s of=%s bs=%d seek=%d conv=notrunc,fsync'
                  % (infile, outfile, sector_size, primary_first))


@defer.inlineCallbacks
def inject_data(    image, key=None, net=None, dns=None, 
                    remove_network_udev=False, 
                    partition=None, execute=None):
    """Injects a ssh key and optionally net data into a disk image.

    it will mount the image as a fully partitioned disk and attempt to inject
    into the specified partition number.

    If partition is not specified it mounts the image as a single partition.

    """
    out, err = execute('sudo losetup --find --show %s' % image)
    if err:
        raise exception.Error('Could not attach image to loopback: %s' % err)
    device = out.strip()
    try:
        if not partition is None:
            # create partition
            out, err = execute('sudo kpartx -a %s' % device)
            if err:
                raise exception.Error('Failed to load partition: %s' % err)
            mapped_device = '/dev/mapper/%sp%s' % (device.split('/')[-1],
                                                   partition)
        else:
            mapped_device = device

        # We can only loopback mount raw images.  If the device isn't there,
        #  it's normally because it's a .vmdk or a .vdi etc
        if not os.path.exists(mapped_device):
            raise exception.Error(
                'Mapped device was not found (we can only inject raw disk images): %s'
                % mapped_device)

        # Configure ext2fs so that it doesn't auto-check every N boots
        out, err = execute('sudo tune2fs -c 0 -i 0 %s' % mapped_device)

        tmpdir = tempfile.mkdtemp()
        try:
            # mount loopback to dir
            out, err = execute(
                    'sudo mount %s %s' % (mapped_device, tmpdir))
            if err:
                raise exception.Error('Failed to mount filesystem: %s' % err)

            try:
                if key:
                    # inject key file
                    _inject_key_into_fs(key, tmpdir, execute=execute)
                if net:
                    _inject_net_into_fs(net, tmpdir, execute=execute)
                if dns:
                    _inject_dns_into_fs(dns, tmpdir, execute=execute)
                if remove_network_udev:
                    _remove_network_udev(tmpdir, execute=execute)

            finally:
                # unmount device
                execute('sudo umount %s' % mapped_device)
        finally:
            # remove temporary directory
            execute('rmdir %s' % tmpdir)
            if not partition is None:
                # remove partitions
                execute('sudo kpartx -d %s' % device)
    finally:
        # remove loopback
        execute('sudo losetup --detach %s' % device)

def _inject_key_into_fs(key, fs, execute=None):
    sshdir = os.path.join(fs, 'root', '.ssh')
    execute('sudo mkdir -p %s' % sshdir) # existing dir doesn't matter
    execute('sudo chown root %s' % sshdir)
    execute('sudo chmod 700 %s' % sshdir)
    keyfile = os.path.join(sshdir, 'authorized_keys')
    execute('sudo tee -a %s' % keyfile, '\n' + key.strip() + '\n')


def _inject_net_into_fs(net, fs, execute=None):
    netfile = os.path.join(fs, 'etc', 'network', 'interfaces')
    execute('sudo tee %s' % netfile, net)

def _inject_dns_into_fs(dns, fs, execute=None):
    dnsfile = os.path.join(fs, 'etc', 'resolv.conf')
    execute('sudo tee %s' % dnsfile, dns)

def _remove_network_udev(fs, execute=None):
    # TODO(justinsb): This is correct for Ubuntu, but might not be right for
    #  other distros.  There is a much bigger discussion to be had about what
    #  we inject and how we inject it.
    rulesfile = os.path.join(fs, 'etc', 'udev', 'rules.d', '70-persistent-net.rules')
    execute('rm -f %s' % rulesfile)