1##########################################################################
2# Copyright (c) 2016, ETH Zurich.
3# All rights reserved.
4#
5# This file is distributed under the terms in the attached LICENSE file.
6# If you do not find this file, copies can be found by writing to:
7# ETH Zurich D-INFK, Universitaetstr 6, CH-8092 Zurich. Attn: Systems Group.
8##########################################################################
9
10import subprocess
11import os
12
13class EFIImage:
14
15    def __init__(self, image, size):
16        """
17        Size in MiB
18        """
19        self._image = image
20        # Size of the disk
21        self._sizeMB= size * 1024 * 1024
22        # block size
23        self._blockSize = 512
24        # size of the disk in blocks
25        self._sizeBlocks= self._sizeMB / self._blockSize
26        # first block of partition
27        self._startBlock = 2048
28
29        # size of partition in blocks
30        self._partSizeBlocks = self._sizeBlocks - self._startBlock
31
32        # calculate byte offset to first partition and format as mformat name
33        self._mformatImage="%s@@%d" % (self._image, self._startBlock*self._blockSize)
34
35        self._dirs = None
36
37    def _cmd(self, command, **kwargs):
38        print(" ".join(command))
39        return subprocess.check_call(command, **kwargs)
40
41
42    def create(self):
43        self._cmd(["dd", "if=/dev/zero", "of=%s" % self._image,
44                  "bs=%d" % self._blockSize, "count=1",
45                  "seek=%d" % (self._sizeBlocks - 1)])
46
47        self._cmd(["/sbin/parted", "-s", self._image, "mktable", "gpt"])
48        self._cmd(["/sbin/parted", "-s", self._image, "mkpart", "primary", "fat32",
49                  "%ds" % self._startBlock, "%ds" % self._partSizeBlocks])
50        self._cmd(["/sbin/parted", "-s", self._image, "align-check", "optimal", "1"])
51        self._cmd(["/sbin/parted", "-s", self._image, "name", "1", "UEFI"])
52
53        self._cmd(["mformat", "-i", self._mformatImage, "-T",
54                  str(self._partSizeBlocks), "-h", "1", "-s", "1"])
55        # mdir fails if the root directory is empty. We create a directory here
56        # to make sure _initDirCache does not fail.
57        self._cmd(["mmd", "-i", self._mformatImage, "dummy"])
58
59        # reset directory cache
60        self._dirs = None
61
62    def _initDirCache(self):
63        if not self._dirs is None:
64            return
65        self._dirs = set()
66        cmd = ["mdir", "-i", self._mformatImage, "-/b"]
67        print(" ".join(cmd))
68        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
69        for line in proc.stdout:
70            if line.endswith("/"):
71                self._dirs.add(line[:-1])
72
73    def _createParentDir(self, dirName):
74        """ Create a parent directory for passed directory name """
75        parentDir = os.path.dirname(dirName)
76        if parentDir is "":
77            return "::"
78        basename = "::/%s" % parentDir
79        if not basename in self._dirs:
80            self._createParentDir(parentDir)
81            self._cmd(["mmd", "-i", self._mformatImage, basename])
82            self._dirs.add(basename)
83        return basename
84
85    def addFile(self, inFile, fileName):
86        if self._dirs is None:
87            self._initDirCache()
88        dirName = self._createParentDir(fileName)
89        targetFile = os.path.join(dirName, os.path.basename(fileName))
90        self._cmd(["mcopy", "-o", "-s", "-i", self._mformatImage, inFile, targetFile])
91
92    def writeFile(self, fileName, contents):
93        if self._dirs is None:
94            self._initDirCache()
95        dirName = self._createParentDir(fileName)
96        cmd = ["mcopy", "-o", "-s", "-i", self._mformatImage, "-", os.path.join(dirName, os.path.basename(fileName))]
97        print(" ".join(cmd))
98        proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
99        proc.communicate(contents)
100        proc.stdin.close()
101