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