java-ebpf

1
#!/usr/bin/env python3
2

3
"""
4
Usage:
5
    python3 find_and_get_kernel.py <version> <destination_path>
6

7
Obtain the linux headers deb package for the given kernel version from
8
https://kernel.ubuntu.com/mainline/v<version>/<arch>/,
9
and place them in <destination_path>/lib/modules/<version>/build.
10

11
Doesn't have any dependencies except for python3 and dpkg-deb.
12
"""
13

14
import os
15
import shutil
16
import zipfile
17
from urllib.request import urlretrieve
18
from pathlib import Path
19
import re
20
import argparse
21

22
BASEDIR = Path(__file__).parent
23
CACHEDIR = BASEDIR / ".cache"
24

25

26
def get_arch() -> str:
27
    import platform
28
    if platform.machine() == "x86_64":
29
        return "amd64"
30
    elif platform.machine() == "aarch64":
31
        return "arm64"
32
    else:
33
        raise Exception("Unknown architecture")
34

35

36
def download_file(url: str, name: str) -> Path:
37
    name = name.replace("/", "_")
38
    os.makedirs(BASEDIR / ".cache", exist_ok=True)
39
    path = CACHEDIR / name
40
    if not path.exists():
41
        print(f"Downloading {url}")
42
        urlretrieve(url, path)
43
    return path
44

45

46
def ci_version_to_mainline_version(version: str) -> str:
47
    if re.match(r"^\d+\.\d+$", version):
48
        return version
49
    if re.match(r"^\d+\.\d+\.\d+$", version):
50
        return version[:3]
51
    raise Exception(f"Unsupported version format: {version}")
52

53

54
def get_deb_url(arch: str, version: str) -> str:
55
    version = ci_version_to_mainline_version(version)
56
    url = f"https://kernel.ubuntu.com/mainline/v{version}/{arch}"
57
    with download_file(url, f"index{arch}{version}").open() as f:
58
        # find <a href="linux-headers-6.6.0-060600-generic_6.6.0-060600.202311151808_amd64.deb">
59
        for line in f:
60
            m = re.search(fr'<a href="linux-headers-.*{arch}.deb">',
61
                          line)
62
            if m:
63
                file = m.group(0).split('"')[1]
64
                return f"https://kernel.ubuntu.com/mainline/v{version}/{arch}/{file}"
65

66

67
def download_deb(arch: str, version: str) -> Path:
68
    url = get_deb_url(arch, version)
69
    return download_file(url, url.split("/")[-1])
70

71

72
def unpack_deb_repack_in_cache(arch: str, version: str) -> Path:
73
    """
74
    Download the linux headers deb package and repackage it into a zip file.
75

76
    The zip file contains `/lib` and `/usr` folders on the top level.
77

78
    Returns: path to the zip file
79
    """
80
    # unpack deb into tmp folder via dpkg-deb
81
    # zip tmp folder and store in .cache/arch-version.zip
82
    # only unpack deb and zip if not already in cache
83
    zip_path = CACHEDIR / f"{arch}-{version}.zip"
84
    if not zip_path.exists():
85
        deb_path = download_deb(arch, version)
86
        tmp_path = CACHEDIR / f"{arch}-{version}"
87
        os.makedirs(tmp_path, exist_ok=True)
88
        os.system(f"dpkg-deb -x {deb_path} {tmp_path}")
89
        zip_dir = CACHEDIR / f"{arch}-{version}"
90
        excluded_folders = []
91
        for folder in tmp_path.glob("lib/modules/*/build/*"):
92
            if folder.name in ["drivers", "scripts"]:
93
                excluded_folders.append(folder.relative_to(tmp_path))
94
            elif folder.name == "arch":
95
                for arch_folder in folder.glob("*"):
96
                    if arch_folder.name not in ["x86", "arm64"]:
97
                        excluded_folders.append(
98
                            arch_folder.relative_to(tmp_path))
99
            elif folder.name == "tools":
100
                for tools_folder in folder.glob("*"):
101
                    if tools_folder.name != "bpf":
102
                        excluded_folders.append(tools_folder.relative_to(tmp_path))
103
        exclude_file = tmp_path / "exclude.txt"
104
        with exclude_file.open("w") as f:
105
            for folder in excluded_folders:
106
                f.write(f"{folder}/\n")
107

108
        os.system(
109
            f"cd {zip_dir}; zip -D -qr {zip_path} {tmp_path.relative_to(zip_dir)} -x@{exclude_file}")
110
        shutil.rmtree(tmp_path)
111
        deb_path.unlink()
112
    return zip_path
113

114

115
def copy_headers_into_dest(arch: str, version: str, dest_root: Path):
116
    """
117
    Download the linux headers and move the "zip.zip:/lib/modules/*/build" folder
118
    to "dest_root/lib/modules/*/build".
119
    """
120
    zip_path = unpack_deb_repack_in_cache(arch, version)
121
    with zipfile.ZipFile(zip_path) as zip_file:
122
        for name in zip_file.namelist():
123
            if not name.startswith("lib/modules"):
124
                continue
125
            inner_name = "/".join(name.split("/")[3:])
126
            if inner_name.startswith("build"):
127
                dest_file = dest_root / "lib" / "modules" / version / inner_name
128
                os.makedirs(dest_file.parent, exist_ok=True)
129
                with dest_file.open("wb") as f:
130
                    f.write(zip_file.read(name))
131

132

133
if __name__ == '__main__':
134
    argparse = argparse.ArgumentParser()
135
    argparse.add_argument("version", help="Kernel version")
136
    argparse.add_argument("destination", help="Destination folder")
137
    args = argparse.parse_args()
138
    copy_headers_into_dest(get_arch(),
139
                           args.version,
140
                           Path(args.destination))
141