├── requirements.txt ├── pyproject.toml ├── LICENSE ├── vendor └── K12.py ├── README.md └── straw.py /requirements.txt: -------------------------------------------------------------------------------- 1 | dataclasses; python_version < '3.7' 2 | pymunge==0.1.3 3 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "straw" 3 | version = "0.1" 4 | readme = "README.md" 5 | keywords = ["slurm"] 6 | license = {text = "BSD 3-Clause License"} 7 | dependencies = [ 8 | 'dataclasses; python_version < "3.7"', 9 | ] 10 | 11 | [project.optional-dependencies] 12 | munge = ["pymunge==0.1.3"] 13 | 14 | [project.scripts] 15 | straw = 'straw:main' 16 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2022, Pablo Llopis 4 | 5 | Redistribution and use in source and binary forms, with or without 6 | modification, are permitted provided that the following conditions are met: 7 | 8 | 1. Redistributions of source code must retain the above copyright notice, this 9 | list of conditions and the following disclaimer. 10 | 11 | 2. Redistributions in binary form must reproduce the above copyright notice, 12 | this list of conditions and the following disclaimer in the documentation 13 | and/or other materials provided with the distribution. 14 | 15 | 3. Neither the name of the copyright holder nor the names of its 16 | contributors may be used to endorse or promote products derived from 17 | this software without specific prior written permission. 18 | 19 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 20 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 21 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 22 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 23 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 24 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 25 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 26 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 27 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 29 | -------------------------------------------------------------------------------- /vendor/K12.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Implementation by Gilles Van Assche, hereby denoted as "the implementer". 3 | # 4 | # For more information, feedback or questions, please refer to our website: 5 | # https://keccak.team/ 6 | # 7 | # To the extent possible under law, the implementer has waived all copyright 8 | # and related or neighboring rights to the source code in this file. 9 | # http://creativecommons.org/publicdomain/zero/1.0/ 10 | 11 | def ROL64(a, n): 12 | return ((a >> (64-(n%64))) + (a << (n%64))) % (1 << 64) 13 | 14 | def KeccakP1600onLanes(lanes, nrRounds): 15 | R = 1 16 | for round in range(24): 17 | if (round + nrRounds >= 24): 18 | # θ 19 | C = [lanes[x][0] ^ lanes[x][1] ^ lanes[x][2] ^ lanes[x][3] ^ lanes[x][4] for x in range(5)] 20 | D = [C[(x+4)%5] ^ ROL64(C[(x+1)%5], 1) for x in range(5)] 21 | lanes = [[lanes[x][y]^D[x] for y in range(5)] for x in range(5)] 22 | # ρ and π 23 | (x, y) = (1, 0) 24 | current = lanes[x][y] 25 | for t in range(24): 26 | (x, y) = (y, (2*x+3*y)%5) 27 | (current, lanes[x][y]) = (lanes[x][y], ROL64(current, (t+1)*(t+2)//2)) 28 | # χ 29 | for y in range(5): 30 | T = [lanes[x][y] for x in range(5)] 31 | for x in range(5): 32 | lanes[x][y] = T[x] ^((~T[(x+1)%5]) & T[(x+2)%5]) 33 | # ι 34 | for j in range(7): 35 | R = ((R << 1) ^ ((R >> 7)*0x71)) % 256 36 | if (R & 2): 37 | lanes[0][0] = lanes[0][0] ^ (1 << ((1<> 7)*0x71)) % 256 41 | return lanes 42 | 43 | def load64(b): 44 | return sum((b[i] << (8*i)) for i in range(8)) 45 | 46 | def store64(a): 47 | return bytearray((a >> (8*i)) % 256 for i in range(8)) 48 | 49 | def KeccakP1600(state, nrRounds): 50 | lanes = [[load64(state[8*(x+5*y):8*(x+5*y)+8]) for y in range(5)] for x in range(5)] 51 | lanes = KeccakP1600onLanes(lanes, nrRounds) 52 | state = bytearray().join([store64(lanes[x][y]) for y in range(5) for x in range(5)]) 53 | return bytearray(state) 54 | 55 | def F(inputBytes, delimitedSuffix, outputByteLen): 56 | outputBytes = bytearray() 57 | state = bytearray([0 for i in range(200)]) 58 | rateInBytes = 1344//8 59 | blockSize = 0 60 | inputOffset = 0 61 | # === Absorb all the input blocks === 62 | while(inputOffset < len(inputBytes)): 63 | blockSize = min(len(inputBytes)-inputOffset, rateInBytes) 64 | for i in range(blockSize): 65 | state[i] = state[i] ^ inputBytes[i+inputOffset] 66 | inputOffset = inputOffset + blockSize 67 | if (blockSize == rateInBytes): 68 | state = KeccakP1600(state, 12) 69 | blockSize = 0 70 | # === Do the padding and switch to the squeezing phase === 71 | state[blockSize] = state[blockSize] ^ delimitedSuffix 72 | if (((delimitedSuffix & 0x80) != 0) and (blockSize == (rateInBytes-1))): 73 | state = KeccakP1600(state, 12) 74 | state[rateInBytes-1] = state[rateInBytes-1] ^ 0x80 75 | state = KeccakP1600(state, 12) 76 | # === Squeeze out all the output blocks === 77 | while(outputByteLen > 0): 78 | blockSize = min(outputByteLen, rateInBytes) 79 | outputBytes = outputBytes + state[0:blockSize] 80 | outputByteLen = outputByteLen - blockSize 81 | if (outputByteLen > 0): 82 | state = KeccakP1600(state, 12) 83 | return outputBytes 84 | 85 | def right_encode(x): 86 | S = bytearray() 87 | while(x > 0): 88 | S = bytearray([x % 256]) + S 89 | x = x//256 90 | S = S + bytearray([len(S)]) 91 | return S 92 | 93 | # inputMessage and customizationString must be of type byte string or byte array 94 | def KangarooTwelve(inputMessage, customizationString, outputByteLen): 95 | B = 8192 96 | c = 256 97 | S = bytearray(inputMessage) + bytearray(customizationString) + right_encode(len(customizationString)) 98 | # === Cut the input string into chunks of B bytes === 99 | n = (len(S)+B-1)//B 100 | Si = [bytearray(S[i*B:(i+1)*B]) for i in range(n)] 101 | if (n == 1): 102 | # === Process the tree with only a final node === 103 | return F(Si[0], 0x07, outputByteLen) 104 | else: 105 | # === Process the tree with kangaroo hopping === 106 | CVi = [F(Si[i+1], 0x0B, c//8) for i in range(n-1)] 107 | NodeStar = Si[0] + bytearray([3,0,0,0,0,0,0,0]) + bytearray().join(CVi) \ 108 | + right_encode(n-1) + b'\xFF\xFF' 109 | return F(NodeStar, 0x06, outputByteLen) 110 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 🥤 Straw for your Slurm beverage! 2 | 3 | Straw is a simple and minimalistic one-shot cli tool that fetches the Slurm config files from a Slurm server running the slurmctld. 4 | It can greatly simplify the deployment of (containerised) environments that interact as clients with Slurm clusters by removing the need 5 | for maintaining munge keys, Slurm config files, as well as slurmd, and munge daemons. 6 | 7 | ## Why Straw? 8 | 9 | In order to create tools and clients that interact with a Slurm cluster, an environment/container usually needs at least the following: 10 | * The Slurm config file(s). 11 | * The Munge authentication tokens. 12 | 13 | Most of the time, this involves: 14 | * Either maintaining a copy of your Slurm config file(s), or running slurmd with configless mode. 15 | * The munged daemon configured with a munge key for authentication 16 | 17 | When containerising tools or clients that interact with the Slurm cluster, it is undesirable having to run 18 | and setup these extra services on each container, and managing the munge.key that is shared in the cluster requires utmost care. 19 | 20 | In fact, when building Slurm client environments (such as containers), due to how Slurm tools have been designed, 21 | you must choose between either a *configless* setup, or a munge *secretless* setup. 22 | 23 | *Configless setup* refers to an environment which is agnostic to Slurm config files (no need update and put these files in your environment). 24 | 25 | *Secretless setup* refers to not needing to share the secret munge key with your environment. 26 | 27 | For instance, you may want to connect some notebook service that is exposed to the internet to your Slurm cluster. In this situation, you might prefer not to keep the Slurm munge key in the public notebook service that's exposed to the wider Internet. 28 | 29 | ``` 30 | D M Z firewall 31 | │ 32 | ┌────────────┐ │ ┌───────────┐ 33 | │ │ │ │ │ 34 | │ Public │ │ │ Private │ 35 | │ notebook │ │ │ Slurm │ 36 | │ service │ │ │ cluster │ 37 | │ │ │ │ │ 38 | └────────────┘ │ └───────────┘ 39 | ``` 40 | 41 | One way to go munge-secretless is to rely on JWT tokens (which arguably are a secret, but where the risk implications are much lower than the munge key). 42 | However, Slurm tools can not use JWT tokens and configless mode simultaneously. 43 | And while it is possible to use JWT in combination with a minimalistic slurm.conf with just the `SlurmctldHost` and `Clustername`, 44 | some commands such as srun refuse to run without more information in the slurm.conf, limiting the usefulness of this approach. 45 | 46 | Straw is a one-shot cli tool that aims to greatly simplify these use cases, and provide increased security for these environments. 47 | 48 | Straw fetches the Slurm config files, optionally using JWT authentication, removing the burden of setting up and running both slurmd and munged, 49 | and allowing *configless* and *munge-secretless* environments to interact with Slurm clusters. 50 | 51 | ## How does it work? 52 | 53 | Straw talks just enough of the Slurm protocol to be able to retrieve the config files, and perform either munge or jwt authentication, 54 | in a way that regular Slurm tools don't. 55 | 56 | Therefore you can create containers that do not contain neither the slurm.conf files, nor munge secrets, nor run additional daemons (such as munge or slurmd). 57 | Just add straw to your container and call it early on during initialisation, ensuring your environment starts after having fetched all Slurm config files. 58 | 59 | ## Requirements 60 | 61 | Running this tool requires python 3. 62 | 63 | If you're using munge to authenticate, you must run this tool as either 64 | the Slurm user, or root. The slurmctld needs to have [configless mode](https://slurm.schedmd.com/configless_slurm.html) enabled as well. 65 | Optionally, for JWT authentication you'll need to enable [JWT support](https://slurm.schedmd.com/jwt.html) in your slurmctld. 66 | 67 | ## Building 68 | 69 | With munge support 70 | ``` 71 | pip install "straw[munge] @ git+https://github.com/pllopis/straw" 72 | ``` 73 | 74 | Without munge support 75 | ``` 76 | pip install "straw @ git+https://github.com/pllopis/straw" 77 | ``` 78 | 79 | ## Usage 80 | 81 | ``` 82 | usage: straw.py [-h] [--auth {munge,jwt}] [-o OUTPUT_DIR] [-v] [-V] [-l] server [server ...] version 83 | 84 | positional arguments: 85 | server slurmctld server in server[:port] notation 86 | version Slurm major version that corresponds to that of the slurmctld server (e.g. 22.05) 87 | 88 | options: 89 | -h, --help show this help message and exit 90 | --auth {munge,jwt} Authentication method (default: jwt) 91 | -o OUTPUT_DIR, --output-dir OUTPUT_DIR 92 | Existing output directory where config files will be saved (default: ./) 93 | -v, --verbose Increase output verbosity. Rrepetitions allowed. (default: None) 94 | -V, --version show program's version number and exit 95 | -l, --list List available protocol versions (default: False) 96 | ``` 97 | 98 | Where auth\_method is either `munge` or `jwt`. The `pymunge` import is conditional on using munge as authentication method, so if yo do not need munge, the library requirement is also not needed. 99 | When using jwt authentication, the token will be grabbed from the `SLURM_JWT` environment variable. 100 | 101 | The Slurm version should include the major release (first two parts), e.g. `22.05`. 102 | It should also match that of the slurmctld server, as this will determine the Slurm protocol version that straw will use to communicate with the slurmctld. 103 | -------------------------------------------------------------------------------- /straw.py: -------------------------------------------------------------------------------- 1 | """ 2 | Straw, the simple tool to suck the config out of your Slurm beverage! 3 | """ 4 | import sys 5 | import os 6 | import socket 7 | import argparse 8 | import re 9 | import logging 10 | 11 | from dataclasses import dataclass, field 12 | from vendor.K12 import KangarooTwelve 13 | from struct import pack, unpack, calcsize 14 | try: 15 | from pymunge import MungeContext, UID_ANY, GID_ANY 16 | except ImportError: 17 | has_munge = False 18 | else: 19 | has_munge = True 20 | 21 | SLURM_PROTOCOL_VERSION = { 22 | '22.05': (38 << 8) | 0, 23 | '21.08': (37 << 8) | 0, 24 | '20.11': (36 << 8) | 0, 25 | } 26 | 27 | HASH_K12 = 2 28 | HASH_K12_LEN = 32 29 | PLUGIN_AUTH_MUNGE = 0x0065 30 | PLUGIN_AUTH_JWT = 0x0066 31 | REQUEST_CONFIG = 0x07df 32 | RESPONSE_CONFIG = REQUEST_CONFIG+1 33 | RESPONSE_SLURM_RC = 8001 34 | NO_VAL = 0xfffffffe 35 | 36 | protocol_version = None 37 | 38 | def list_protocol_versions(): 39 | for ver in SLURM_PROTOCOL_VERSION.keys(): 40 | print(ver) 41 | 42 | @dataclass 43 | class Header: 44 | # default values for packing the request_config msg 45 | protocol_version: int = 0 46 | flags: int = 0 47 | msg_type: int = REQUEST_CONFIG 48 | body_length: int = 4 49 | forward_cnt: int = 0 50 | ret_cnt: int = 0 51 | address_ss_family: int = 0 52 | 53 | 54 | @dataclass 55 | class Auth: 56 | """ 57 | To use with munge, use Auth(plugin_id=PLUGIN_AUTH_MUNGE). Munge cred will be generated automatically. 58 | To use with JWT, use Auth(cred=jwt_token, plugin_id=PLUGIN_AUTH_JWT). 59 | """ 60 | plugin_id: int = None 61 | cred: str = None 62 | 63 | class Body: 64 | pass 65 | 66 | @dataclass 67 | class RequestConfigBody(Body): 68 | req_flags: int = 0x001 69 | 70 | def pack(self): 71 | return pack('!I', self.req_flags) 72 | 73 | @dataclass 74 | class ResponseConfigBody(Body): 75 | config_files: list[tuple] = field(default_factory=list) # filename, content 76 | spool_dir: str = None 77 | 78 | def pack(self): 79 | raise NotImplementedError('Responses are only unpacked, not packed') 80 | 81 | 82 | @dataclass 83 | class SlurmMessage: 84 | """ 85 | A Slurm message consists of `header`, `auth`, and `body` payloads, in that order. 86 | When packing a message, because they depend on each other (e.g. auth may need to create a hash of the body), 87 | first the separate instances are created. Then, they can be pack()ed. 88 | When unpacking a message, the binary `data` will be "walked", starting at offset 0, 89 | unpacking header, auth, and body, also in that order. 90 | """ 91 | data: bytes = None 92 | 93 | header: Header = Header() 94 | auth: Auth = Header() 95 | body: Body = None 96 | 97 | _unpack_offset: int = 0 98 | 99 | def _unpack(self, fmt): 100 | """ 101 | Akin to struct.unpack, but keep track of which bytes we've unpacked from self.data (in self._unpack_offset). 102 | """ 103 | unpack_sz = calcsize(fmt) 104 | res = unpack(fmt, self.data[self._unpack_offset:self._unpack_offset+unpack_sz]) 105 | self._unpack_offset += unpack_sz 106 | return res 107 | 108 | def _unpackstr(self): 109 | """ 110 | Unpacks a string <\0>, where len of str is given by the first uint32_t size. 111 | Keep track of the offset in self.data 112 | """ 113 | str_len, = self._unpack('!I') 114 | if (str_len > 0): 115 | # str_len accounts for trailing '\0' 116 | s = self.data[self._unpack_offset:self._unpack_offset+str_len-1] 117 | self._unpack_offset += str_len 118 | return s 119 | else: 120 | return None 121 | 122 | def unpack_header(self): 123 | self.header.protocol_version, = self._unpack('!H') 124 | if self.header.protocol_version != protocol_version: 125 | sys.exit(f'Protocol version response from server ({self.header.protocol_version}) is different to the protocol version requested ({protocol_version}).') 126 | if protocol_version >= SLURM_PROTOCOL_VERSION['22.05']: 127 | _, self.header.msg_type, self.header.body_length, _, _, _ = self._unpack('!HHIHHH') 128 | elif protocol_version >= SLURM_PROTOCOL_VERSION['20.11']: 129 | _, _, self.header.msg_type, self.header.body_length, _, _, _ = self._unpack('!HHHIHHH') 130 | 131 | def unpack_auth(self): 132 | self.auth.plugin_id, = self._unpack('!I') 133 | if self.auth.plugin_id == PLUGIN_AUTH_MUNGE: 134 | self.auth.cred = self._unpackstr() 135 | logging.debug(f'Munge cred: {self.auth.cred}') 136 | if self.auth.plugin_id == PLUGIN_AUTH_JWT: 137 | token = self._unpackstr() 138 | user = self._unpackstr() 139 | 140 | def _unpack_list(self): 141 | lst = [] 142 | config_file_count, = self._unpack('!I') 143 | if config_file_count != NO_VAL: 144 | for i in range(0, config_file_count): 145 | file_exists, = self._unpack('B') 146 | filename = self._unpackstr() 147 | if filename: 148 | filename = filename.decode('utf-8') 149 | content = self._unpackstr() 150 | if content: 151 | content = content.decode('utf-8') 152 | logging.debug(f'filename: {filename}, exists: {bool(file_exists)}') 153 | if file_exists: 154 | lst.append((filename, content)) 155 | return lst 156 | 157 | def unpack_body(self): 158 | if self.header.msg_type == RESPONSE_SLURM_RC: 159 | # Uh oh, error! 160 | rc, = self._unpack('!I') 161 | if rc == 2010: 162 | logging.error('Maybe you did not run as Slurm user or root (required for munge auth)?') 163 | sys.exit(f'We got a reply with errno: {rc}') 164 | elif self.header.msg_type == RESPONSE_CONFIG: 165 | if self.header.protocol_version >= SLURM_PROTOCOL_VERSION['21.08']: 166 | self.body = ResponseConfigBody() 167 | logging.debug(f'Got a response body of length {self.header.body_length}:') 168 | self.body.config_files = self._unpack_list() 169 | self.body.spool_dir = self._unpackstr().decode('utf-8') 170 | elif self.header.protocol_version >= SLURM_PROTOCOL_VERSION['20.11']: 171 | raise NotImplementedError('Fetching config from Slurm 21.08 > version >= 20.11 not yet implemented') 172 | else: 173 | raise Exception(f'Server replied with unsupported protocol version: {self.header.protocol_version}') 174 | 175 | 176 | def unpack(self): 177 | self._unpack_offset = 0 178 | self.unpack_header() 179 | self.unpack_auth() 180 | self.unpack_body() 181 | 182 | def _get_munge_cred(self): 183 | custom_string = pack('!H', REQUEST_CONFIG) 184 | body = self.body.pack() 185 | logging.debug('custom str:') 186 | logging.debug(hexdump(custom_string)) 187 | logging.debug('hash input:') 188 | logging.debug(hexdump(body)) 189 | slurm_hash = pack('B32s', HASH_K12, bytes(KangarooTwelve(body, custom_string, HASH_K12_LEN))) 190 | logging.debug('raw hash:') 191 | logging.debug(hexdump(slurm_hash)) 192 | with MungeContext() as ctx: 193 | ctx.uid_restriction = UID_ANY 194 | ctx.gid_restriction = GID_ANY 195 | self.auth.cred = ctx.encode(slurm_hash) 196 | return self.auth.cred 197 | 198 | def pack(self): 199 | self.data = self.pack_header() + \ 200 | self.pack_auth() + \ 201 | self.pack_body() 202 | logging.debug(f'Full message: ({len(self.data)})') 203 | logging.debug(hexdump(self.data)) 204 | return self.data 205 | 206 | def pack_header(self): 207 | if self.header.protocol_version >= SLURM_PROTOCOL_VERSION['22.05']: 208 | header = pack('!HHHIHHH', 209 | self.header.protocol_version, 210 | self.header.flags, 211 | self.header.msg_type, 212 | self.header.body_length, 213 | self.header.forward_cnt, 214 | self.header.ret_cnt, 215 | self.header.address_ss_family) 216 | elif self.header.protocol_version >= SLURM_PROTOCOL_VERSION['20.11']: 217 | header = pack('!HHHHIHHH', 218 | self.header.protocol_version, 219 | self.header.flags, 220 | 0, 221 | self.header.msg_type, 222 | self.header.body_length, 223 | self.header.forward_cnt, 224 | self.header.ret_cnt, 225 | self.header.address_ss_family) 226 | logging.debug(f'Header ({len(header)}):') 227 | logging.debug(hexdump(header)) 228 | return header 229 | 230 | def pack_auth(self): 231 | if self.auth.plugin_id == PLUGIN_AUTH_MUNGE: 232 | try: 233 | self.auth.cred = self._get_munge_cred() 234 | except Exception as err: 235 | sys.exit(f'Failed to generate munge credential:\n{err}') 236 | auth = pack('!II', self.auth.plugin_id, len(self.auth.cred)+1) + self.auth.cred + b'\x00' 237 | elif self.auth.plugin_id == PLUGIN_AUTH_JWT: 238 | # packstr(token) + packstr(NULL) 239 | auth = pack('!II', self.auth.plugin_id, len(self.auth.cred)+1) + bytes(self.auth.cred, 'utf-8') + b'\x00' + b'\x00\x00\x00\x00' 240 | logging.debug(f'Auth: ({len(auth)})') 241 | logging.debug(hexdump(auth)) 242 | return auth 243 | 244 | def pack_body(self): 245 | body = self.body.pack() 246 | logging.debug(f'Body: ({len(body)})') 247 | logging.debug(hexdump(body)) 248 | return body 249 | 250 | def RequestConfigMsg(self, protocol_version, auth_method): 251 | self.header = Header(protocol_version=protocol_version) 252 | self.body = RequestConfigBody() 253 | logging.info(f'Using authentication method: {auth_method}') 254 | if auth_method == 'jwt': 255 | try: 256 | token = os.environ['SLURM_JWT'] 257 | except: 258 | sys.exit('Auth method jwt requested but SLURM_JWT undefined') 259 | self.auth = Auth(cred=token, plugin_id=PLUGIN_AUTH_JWT) 260 | else: 261 | if not has_munge: 262 | sys.exit('Auth method munge requested, but pymunge not available') 263 | self.auth = Auth(plugin_id=PLUGIN_AUTH_MUNGE) 264 | payload = self.pack() 265 | return payload 266 | 267 | 268 | 269 | def hexdump(data): 270 | # Make sure the input data is a bytestring 271 | if not isinstance(data, bytes): 272 | raise TypeError("hexdump() argument must be a bytestring") 273 | 274 | # Initialize variables 275 | addr = 0 276 | lines = [] 277 | 278 | # Loop over the data in blocks of 16 bytes 279 | for i in range(0, len(data), 16): 280 | # Get the current block of data 281 | block = data[i:i+16] 282 | 283 | # Compute the hexadecimal representation of the block 284 | hexstr = " ".join(f"{b:02x}" for b in block) 285 | 286 | # Compute the ASCII representation of the block, using '.' for non-printable characters 287 | asciistr = "".join(chr(b) if 32 <= b < 127 else "." for b in block) 288 | 289 | # Add the address, hexadecimal representation, and ASCII representation to the list of lines 290 | lines.append(f"{addr:08x} {hexstr:47} {asciistr}") 291 | 292 | # Increment the address 293 | addr += 16 294 | 295 | # Return the list of lines as a string, separated by newlines 296 | return "\n".join(lines) 297 | 298 | class StrawConnectionError(Exception): 299 | pass 300 | 301 | def send_recv(server, payload): 302 | def parse_server(server): 303 | """"Parse server[:port]""" 304 | vals = server.split(':') 305 | if len(vals) > 1: 306 | return vals[0], vals[1] 307 | else: 308 | return vals[0], 6817 309 | 310 | payload_msg = pack('!I', len(payload)) + payload 311 | 312 | host, port = parse_server(server) 313 | logging.info(f'Trying {host}:{port}...') 314 | try: 315 | s = socket.create_connection((host, port)) 316 | s.sendall(payload_msg) 317 | except Exception as ex: 318 | logging.error(ex) 319 | raise StrawConnectionError() 320 | 321 | with s, s.makefile(mode='rb') as sfile: 322 | recv_len = sfile.read(4) 323 | logging.debug(f'recvd {len(recv_len)} bytes') 324 | resp_len = int(unpack('!I', recv_len)[0]) 325 | logging.debug(f'Read a message of length {resp_len}') 326 | response = bytes(sfile.read(resp_len)) 327 | return response 328 | 329 | def parse_msg(data): 330 | msg = SlurmMessage(data=data) 331 | msg.unpack() 332 | return msg 333 | 334 | def save_config(msg, output_dir): 335 | for file in msg.body.config_files: 336 | filepath = f'{output_dir}/{file[0]}' 337 | content = str(file[1]) 338 | try: 339 | with open(filepath, 'w') as f: 340 | f.write(content) 341 | except Exception as err: 342 | logging.error(f'Unable to write {filepath}: {err}') 343 | 344 | def fetch_config(servers, auth, output_dir='./'): 345 | logging.debug(f'Using protocol version: {protocol_version}') 346 | payload = SlurmMessage().RequestConfigMsg(protocol_version=protocol_version, auth_method=auth) 347 | response_msg = None 348 | for server in servers: 349 | try: 350 | response_msg = send_recv(server, payload) 351 | except StrawConnectionError as err: 352 | logging.info(err) 353 | logging.error('Connection error. Retrying with next server.') 354 | else: 355 | # Only bother retrying further servers for connection errors 356 | break 357 | 358 | if not response_msg: 359 | sys.exit('Unable to connect and no more servers to try') 360 | 361 | logging.debug(hexdump(response_msg)) 362 | slurm_msg = parse_msg(response_msg) 363 | if slurm_msg.body.spool_dir: 364 | logging.info(f'SlurmdSpoolDir={slurm_msg.body.spool_dir}') 365 | save_config(slurm_msg, output_dir) 366 | 367 | def parse_args(): 368 | def major_version_match(arg): 369 | if not re.fullmatch(r'[0-9]{2,}\.[0-9]+', arg): 370 | raise ValueError('Slurm major version must be specified (e.g. 22.05)') 371 | return arg 372 | 373 | # First parser just for listing protocol versions 374 | list_parser = argparse.ArgumentParser(add_help=False) 375 | list_parser.add_argument('-l', '--list', action='store_true') 376 | list_versions = False 377 | try: 378 | args, _ = list_parser.parse_known_args() 379 | if args.list: 380 | list_versions = True 381 | except: 382 | pass 383 | 384 | if list_versions: 385 | list_protocol_versions() 386 | sys.exit(0) 387 | 388 | # Main parser 389 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) 390 | parser.add_argument('server', type=str, nargs='+', 391 | help='slurmctld server in server[:port] notation') 392 | parser.add_argument('version', type=major_version_match, 393 | help='Slurm major version that corresponds to that of the slurmctld server (e.g. 22.05)') 394 | parser.add_argument('--auth', choices=['munge', 'jwt'], default='jwt', 395 | help='Authentication method') 396 | parser.add_argument('-o', '--output-dir', default='./', 397 | help='Existing output directory where config files will be saved') 398 | parser.add_argument('-v', '--verbose', action='count', 399 | help='Increase output verbosity. Rrepetitions allowed.') 400 | parser.add_argument('-V', '--version', action='version', version='%(prog)s 0.1') 401 | parser.add_argument('-l', '--list', action='store_true', 402 | help='List available protocol versions') 403 | return parser.parse_args() 404 | 405 | def main(): 406 | args = parse_args() 407 | 408 | if not args.verbose: 409 | # default logging level 410 | loglevel = logging.ERROR 411 | elif args.verbose >= 2: 412 | loglevel = logging.DEBUG 413 | elif args.verbose >= 1: 414 | loglevel = logging.INFO 415 | logging.getLogger().setLevel(loglevel) 416 | logging.basicConfig(format='%(message)s', datefmt='%m/%d/%Y %H:%M:%S %Z') 417 | 418 | logging.debug(repr(args)) 419 | 420 | global protocol_version 421 | protocol_version = SLURM_PROTOCOL_VERSION[args.version] 422 | fetch_config(args.server, args.auth, args.output_dir) 423 | 424 | if __name__ == '__main__': 425 | main() 426 | --------------------------------------------------------------------------------