1#!/usr/bin/env python
2
3##########################################################################
4# Copyright (c) 2016, 2019 ETH Zurich.
5# All rights reserved.
6#
7# This file is distributed under the terms in the attached LICENSE file.
8# If you do not find this file, copies can be found by writing to:
9# ETH Zurich D-INFK, Universitaetstr 6, CH-8092 Zurich. Attn: Systems Group.
10##########################################################################
11
12import subprocess
13import os, argparse, re
14
15class EFIImage:
16
17    def __init__(self, image, size):
18        """
19        Size in MiB
20        """
21        self._image = image
22        # Size of the disk
23        self._sizeMB= size * 1024 * 1024
24        # block size
25        self._blockSize = 512
26        # size of the disk in blocks
27        self._sizeBlocks= self._sizeMB / self._blockSize
28        # first block of partition
29        self._startBlock = 2048
30
31        # size of partition in blocks
32        self._partSizeBlocks = self._sizeBlocks - self._startBlock
33
34        # calculate byte offset to first partition and format as mformat name
35        self._mformatImage="%s@@%d" % (self._image, self._startBlock*self._blockSize)
36
37        self._dirs = None
38
39    def _cmd(self, command, **kwargs):
40        print(" ".join(command))
41        return subprocess.check_call(command, **kwargs)
42
43
44    def create(self):
45        self._cmd(["dd", "if=/dev/zero", "of=%s" % self._image,
46                  "bs=%d" % self._blockSize, "count=1",
47                  "seek=%d" % (self._sizeBlocks - 1)])
48
49        self._cmd(["/sbin/parted", "-s", self._image, "mktable", "gpt"])
50        self._cmd(["/sbin/parted", "-s", self._image, "mkpart", "primary", "fat32",
51                  "%ds" % self._startBlock, "%ds" % self._partSizeBlocks])
52        self._cmd(["/sbin/parted", "-s", self._image, "align-check", "optimal", "1"])
53        self._cmd(["/sbin/parted", "-s", self._image, "name", "1", "UEFI"])
54
55        self._cmd(["mformat", "-i", self._mformatImage, "-T",
56                  str(self._partSizeBlocks), "-h", "1", "-s", "1"])
57        # mdir fails if the root directory is empty. We create a directory here
58        # to make sure _initDirCache does not fail.
59        self._cmd(["mmd", "-i", self._mformatImage, "dummy"])
60
61        # reset directory cache
62        self._dirs = None
63
64    def _initDirCache(self):
65        if not self._dirs is None:
66            return
67        self._dirs = set()
68        cmd = ["mdir", "-i", self._mformatImage, "-/b"]
69        print(" ".join(cmd))
70        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
71        for line in proc.stdout:
72            if line.endswith("/"):
73                self._dirs.add(line[:-1])
74
75    def _createParentDir(self, dirName):
76        """ Create a parent directory for passed directory name """
77        parentDir = os.path.dirname(dirName)
78        if parentDir is "":
79            return "::"
80        basename = "::/%s" % parentDir
81        if not basename in self._dirs:
82            self._createParentDir(parentDir)
83            self._cmd(["mmd", "-i", self._mformatImage, basename])
84            self._dirs.add(basename)
85        return basename
86
87    def addFile(self, inFile, fileName):
88        if self._dirs is None:
89            self._initDirCache()
90        dirName = self._createParentDir(fileName)
91        targetFile = os.path.join(dirName, os.path.basename(fileName))
92        self._cmd(["mcopy", "-o", "-s", "-i", self._mformatImage, inFile, targetFile])
93
94    def writeFile(self, fileName, contents):
95        if self._dirs is None:
96            self._initDirCache()
97        dirName = self._createParentDir(fileName)
98        cmd = ["mcopy", "-o", "-s", "-i", self._mformatImage, "-", os.path.join(dirName, os.path.basename(fileName))]
99        print(" ".join(cmd))
100        proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
101        proc.communicate(contents)
102        proc.stdin.close()
103
104
105#
106# Command line and interface functions
107#
108
109def build_bf_efi_img(img_file, build_dir, modules, menulst, hagfish):
110    hagfish_default = os.path.abspath(os.path.join(
111        os.path.dirname(os.path.realpath(__file__)),
112        "..","hagfish","Hagfish.efi"))
113    if hagfish is None:
114        hagfish = hagfish_default
115
116    efi = EFIImage(img_file, 200)
117    efi.create()
118    for module in modules:
119        if module[0] == "/":
120            module = module[1:]
121        efi.addFile(os.path.join(build_dir, module), module)
122    efi.writeFile("startup.nsh", "Hagfish.efi hagfish.cfg")
123    efi.addFile(hagfish, "Hagfish.efi")
124    efi.addFile(menulst, "hagfish.cfg")
125
126def parse_menu_lst(menulst):
127    with open(menulst) as f:
128        res = []
129        for line in f.readlines():
130            line = line.strip()
131            if line=="" or line[0] == "#": continue
132            parts = re.split("\s+", line, 2)
133            if len(parts) < 1:
134                print("Ignoring line: %s" % line)
135                continue
136            if parts[0] in ["stack"]:
137                # no module
138                continue
139            res.append(parts[1])
140
141    return res
142
143def parse_args():
144    p = argparse.ArgumentParser(
145        description='Build EFI image for ARMv8')
146    p.add_argument('menulst', help='the menu.lst')
147    p.add_argument('build', help='build directory')
148    p.add_argument('file', help='output image')
149    p.add_argument('--hagfish', help='hagfish location', default=None)
150    return p.parse_args()
151
152if __name__ == "__main__":
153    print("Barrelfish EFI Image Builder")
154    args = parse_args()
155    print("Image File: %s" % args.file)
156    print("menu.lst: %s" % args.menulst)
157    modules = parse_menu_lst(args.menulst)
158    print("Modules: %s" % ",".join(modules))
159    build_bf_efi_img(args.file, args.build, modules, args.menulst, args.hagfish)
160