Coverage for /Users/buh/.pyenv/versions/3.12.2/envs/pii/lib/python3.12/site-packages/es_pii_tool/helpers/utils.py: 72%

1"""Helper Functions"""

2

3 import typing as t

4 import logging

5 import json

6 from inspect import stack

7 from datetime import datetime, timezone

8 import re

9 from elasticsearch8.exceptions import NotFoundError

10 from es_client.exceptions import ConfigurationError as esc_ConfigError

11 from es_client.helpers.schemacheck import SchemaCheck

12 from es_client.helpers.utils import get_yaml

13 from es_wait.exceptions import IlmWaitError

14 import es_pii_tool.exceptions as e

15 from es_pii_tool.defaults import PHASES, redaction_schema

16

17 if t.TYPE_CHECKING:

18 from dotmap import DotMap # type: ignore

19 from voluptuous import Schema

20 from elasticsearch8 import Elasticsearch

21 from es_pii_tool.job import Job

22 from es_pii_tool.task import Task

23

24 logger = logging.getLogger(__name__)

25

26

27 def build_script(message: str, fields: t.Sequence[str]) -> t.Dict[str, str]:

28 """

29 Build a painless script for redacting fields by way of an update_by_query operation

30

31 :param message: The text to put in place of whatever is in a field

32 :param fields: The list of field names to act on

33

34 :type message: str

35 :type fields: list

36

37 :rtype: dict

38 :returns: A dictionary of ``{"source": (assembled message), "lang": "painless"}``

39 """

40 msg = ""

41 for field in fields:

42 msg += f"ctx._source.{field} = '{message}'; "

43 script = {"source": msg, "lang": "painless"}

44 logger.debug('script = %s', script)

45 return script

46

47

48 def check_dotted_fields(result: t.Dict, field: str, message: str) -> bool:

49 """Iterate through dotted fields to ensure success

50

51 :param result: The search result object

52 :param field: The field with dotted notation

53

54 :type result: dict

55 :type field: str

56

57 :returns: Success (``True``) or Failure (``False``)

58 :rtype: bool

59 """

60 success = False

61 logger.debug('Dotted field detected: (%s) ...', field)

62 fielder = result['hits']['hits'][0]['_source']

63 iterations = len(field.split('.'))

64 counter = 1

65 for key in field.split('.'):

66 # This should recursively look for each subkey

67 if key in fielder:

68 fielder = fielder[key]

69 else:

70 break

71 if counter == iterations:

72 if fielder == message:

73 success = True

74 counter += 1

75 return success

76

77

78 def check_fields(result: t.Dict, job_config: t.Dict) -> bool:

79 """Check document fields in result to ensure success

80

81 :param result: The search result object

82 :param job_config: The configuration settings for this job

83

84 :type result: dict

85 :type job_config: dict

86

87 :returns: Success (``True``) or Failure (``False``)

88 :rtype: bool

89 """

90 complete = True

91 hit = result['hits']['hits'][0]['_source']

92 for field in job_config['fields']:

93 success = False

94 if len(field.split('.')) > 1:

95 success = check_dotted_fields(result, field, job_config['message'])

96

97 elif field in hit:

98 if hit[field] == job_config['message']:

99 success = True

100

101 else:

102 logger.warning("Field %s not present in document", field)

103 # Don't need to report the expected fail 2x, so we break the loop here

104 break

105

106 if success:

107 logger.info("Field %s is redacted correctly", field)

108 else:

109 # A single failure is enough to make it a complete failure.

110 complete = False

111 logger.error("Field %s is not redacted correctly", field)

112 return complete

113

114

115 def chunk_index_list(indices: t.Sequence[str]) -> t.Sequence[t.Sequence[str]]:

116 """

117 This utility chunks very large index lists into 3KB chunks.

118 It measures the size as a csv string, then converts back into a list for the return

119 value.

120

121 :param indices: The list of indices

122

123 :type indices: list

124

125 :returns: A list of lists (each a piece of the original ``indices``)

126 :rtype: list

127 """

128 chunks = []

129 chunk = ""

130 for index in indices:

131 if len(chunk) < 3072:

132 if not chunk:

133 chunk = index

134 else:

135 chunk += "," + index

136 else:

137 chunks.append(chunk.split(','))

138 chunk = index

139 chunks.append(chunk.split(','))

140 return chunks

141

142

143 def configure_ilm_policy(task: 'Task', data: 'DotMap') -> None:

144 """

145 Prune phases we've already passed.

146

147 If only_expunge_deletes is True in the job config, set any force_merge_index

148 actions to False.

149 """

150 # Copy the existing policy to a new spot

151 data.new.ilmpolicy = data.ilm.lifecycle.policy

152

153 # Prune phases from existing ILM policy we've already surpassed

154 for phase in list(data.new.ilmpolicy.phases.toDict().keys()):

155 if PHASES.index(data.ilm.explain.phase) > PHASES.index(phase):

156 del data.new.ilmpolicy.phases[phase]

157

158 # Figure out if we're doing force merge

159 fmerge = True

160 if 'forcemerge' in task.job.config:

161 fmkwargs = task.job.config['forcemerge']

162 if 'only_expunge_deletes' in fmkwargs and fmkwargs['only_expunge_deletes']:

163 fmerge = False

164 else:

165 fmerge = False

166

167 # Loop through the remaining phases and set 'force_merge_index': False

168 # to the cold or frozen actions.

169

170 for phase in data.new.ilmpolicy.phases:

171 if phase not in ['cold', 'frozen']:

172 continue

173 if 'searchable_snapshot' in data.new.ilmpolicy.phases[phase].actions:

174 data.new.ilmpolicy.phases[

175 phase

176 ].actions.searchable_snapshot.force_merge_index = fmerge

177

178

179 def end_it(obj: t.Union['Job', 'Task'], success: bool) -> None:

180 """Close out the object here to avoid code repetition"""

181 # Record task success or fail here for THIS task_id

182 # Each index in per_index has its own status tracker

183 if not success:

184 err = True

185 log = 'Check application logs for detailed report'

186 else:

187 err = False

188 log = 'DONE'

189 obj.end(completed=success, errors=err, logmsg=log)

190

191

192 def exception_msgmaker(exc: t.Union[e.MissingIndex, e.BadClientResult]) -> str:

193 """Most of the messages here are similar enough to warrant a single function"""

194 upstream = (

195 f'The upstream exception type was {type(exc.upstream).__name__}, '

196 f'with error message: {exc.upstream.args[0]}'

197 )

198 if isinstance(exc, e.MissingIndex):

199 msg = (

200 f'Exception raised because index {exc.missing} was not found. '

201 f'{upstream}'

202 )

203 elif isinstance(exc, e.BadClientResult):

204 msg = (

205 f'Exception raised because of a bad or unexpected response or result '

206 f'from the Elasticsearch cluster. {upstream}'

207 )

208 return msg

209

210

211 def get_alias_actions(oldidx: str, newidx: str, aliases: t.Dict) -> t.Sequence:

212 """

213 :param oldidx: The old index name

214 :param newidx: The new index name

215 :param aliases: The aliases

216

217 :type oldidx: str

218 :type newidx: str

219 :type aliases: dict

220

221 :returns: A list of actions suitable for

222 :py:meth:`~.elasticsearch.client.IndicesClient.update_aliases` ``actions``

223 kwarg.

224 :rtype: list

225 """

226 actions = []

227 for alias in aliases.keys():

228 actions.append({'remove': {'index': oldidx, 'alias': alias}})

229 actions.append({'add': {'index': newidx, 'alias': alias}})

230 return actions

231

232

233 def get_field_matches(config: t.Dict, result: t.Dict) -> int:

234 """Count docs which have the expected fields

235

236 :param config: The config from the YAML file

237 :param result: The query result dict

238

239 :type config: dict

240 :type result: dict

241

242 :returns: The count of docs in ``result`` which have the identified fields

243 :rtype: int

244 """

245

246 logger.debug('Extracting doc hit count from result')

247 doc_count = result['hits']['total']['value']

248 for element in range(0, result['hits']['total']['value']):

249 for field in config['fields']:

250 if len(field.split('.')) > 1:

251 logger.debug('Dotted field "%s" detected...', field)

252 fielder = result['hits']['hits'][element]['_source']

253 for key in field.split('.'):

254 # This should recursively look for each subkey

255 if key in fielder:

256 fielder = fielder[key]

257 else:

258 doc_count -= 1

259 break

260 elif field not in list(result['hits']['hits'][element]['_source'].keys()):

261 logger.debug('Fieldname "%s" NOT detected...', field)

262 doc_count -= 1

263 else:

264 logger.debug('Root-level fieldname "%s" detected...', field)

265 return doc_count

266

267

268 def get_fname() -> str:

269 """Return the name of the calling function"""

270 return stack()[1].function

271

272

273 def get_inc_version(name: str) -> int:

274 """Extract the incrementing version value from the end of name

275

276 :param name: The name

277

278 :type name: str

279

280 :returns: The integer value of the current index revision, or 0 if no version

281 :rtype: int

282 """

283 # Anchor the end as 3 dashes, a v, and 3 digits, e.g. ---v001

284 match = re.search(r'^.*---v(\d{3})$', name)

285 if match:

286 return int(match.group(1))

287 return 0

288

289

290 def get_redactions(file: str = '', data: t.Union[t.Dict, None] = None) -> 'Schema':

291 """

292 Return valid dictionary of redactions from either ``file`` or from ``data``

293 after checking Schema

294

295 :param file: YAML file with redactions to check

296 :param data: Configuration data in dictinoary format

297

298 :type file: str

299 :type data: dict

300

301 :rtype: dict

302 :returns: Redactions configuration data

303 """

304 if data is None:

305 data = {}

306 logger.debug('Getting redactions data...')

307 if file:

308 try:

309 config = get_yaml(file)

310 except esc_ConfigError as exc:

311 msg = f'Unable to read and/or parse YAML REDACTIONS_FILE: {file} Exiting.'

312 logger.critical(msg)

313 raise e.ConfigError(msg, exc)

314 elif data:

315 config = data

316 else:

317 raise e.FatalError('No configuration file or dictionary provided.', Exception())

318 return SchemaCheck(

319 config, redaction_schema(), 'Redaction Configuration', 'redactions'

320 ).result()

321

322

323 def now_iso8601() -> str:

324 """

325 :returns: An ISO8601 timestamp based on datetime.now

326 """

327 # Because Python 3.12 now requires non-naive timezone declarations, we must change.

328 #

329 # ## Example:

330 # ## The new way:

331 # ## datetime.now(timezone.utc).isoformat()

332 # ## Result: 2024-04-16T16:00:00+00:00

333 # ## End Example

334 #

335 # Note that the +00:00 is appended now where we affirmatively declare the UTC

336 # timezone

337 #

338 # As a result, we will use this function to prune away the timezone if it is +00:00

339 # and replace it with Z, which is shorter Zulu notation for UTC (per Elasticsearch)

340 #

341 # We are MANUALLY, FORCEFULLY declaring timezone.utc, so it should ALWAYS be +00:00,

342 # but could in theory sometime show up as a Z, so we test for that.

343

344 parts = datetime.now(timezone.utc).isoformat().split('+')

345 if len(parts) == 1:

346 if parts[0][-1] == 'Z':

347 return parts[0] # Our ISO8601 already ends with a Z for Zulu/UTC time

348 return f'{parts[0]}Z' # It doesn't end with a Z so we put one there

349 if parts[1] == '00:00':

350 return f'{parts[0]}Z' # It doesn't end with a Z so we put one there

351 return f'{parts[0]}+{parts[1]}' # Fallback publishes the +TZ, whatever that was

352

353

354 def config_fieldmap(

355 rw_val: t.Literal['read', 'write'],

356 key: t.Literal[

357 'pattern',

358 'query',

359 'fields',

360 'message',

361 'expected_docs',

362 'restore_settings',

363 'delete',

364 ],

365) -> t.Union[str, int, object]:

366 """

367 Return the function from this function/key map

368 """

369 which = {

370 'read': {

371 'pattern': json.loads,

372 'query': json.loads,

373 'fields': json.loads,

374 'message': str,

375 'expected_docs': int,

376 'restore_settings': json.loads,

377 'delete': str,

378 },

379 'write': {

380 'pattern': json.dumps,

381 'query': json.dumps,

382 'fields': json.dumps,

383 'message': str,

384 'expected_docs': int,

385 'restore_settings': json.dumps,

386 'delete': str,

387 },

388 }

389 return which[rw_val][key]

390

391

392 def parse_job_config(config: t.Dict, behavior: t.Literal['read', 'write']) -> t.Dict:

393 """Parse raw config from the index.

394

395 Several fields are JSON escaped, so we need to fix it to put it in a dict.

396

397 :param config: The raw config data

398 :param behavior: ``read`` or ``write``

399

400 :type config: dict

401 :type behavior: str

402

403 :rtype: dict

404

405 :returns: JSON-(de)sanitized configuration dict

406 """

407 fields = [

408 'pattern',

409 'query',

410 'fields',

411 'message',

412 'expected_docs',

413 'restore_settings',

414 'delete',

415 ]

416 doc = {}

417 for field in fields:

418 if field in config:

419 func = config_fieldmap(behavior, field) # type: ignore

420 doc[field] = func(config[field]) # type: ignore

421 return doc

422

423

424 def strip_ilm_name(name: str) -> str:

425 """

426 Strip leading ``pii-tool-``, and trailing ``---v000`` from ``name``

427

428 :param name: The ILM lifecycle name

429

430 :type name: str

431

432 :returns: The "cleaned up" and stripped ILM name

433 :rtype: str

434 """

435 retval = name.replace('pii-tool-', '')

436 # Anchor the end as 3 dashes, a v, and 3 digits, e.g. ---v001

437 match = re.search(r'^(.*)---v\d{3}$', retval)

438 if match:

439 retval = match.group(1)

440 return retval

441

442

443 def strip_index_name(name: str) -> str:

444 """

445 Strip ``partial-``, ``restored-``, ``redacted-``, and trailing ``---v000`` from

446 ``name``

447

448 :param name: The index name

449

450 :type name: str

451

452 :returns: The "cleaned up" and stripped index name

453 :rtype: str

454 """

455 retval = name.replace('partial-', '')

456 retval = retval.replace('restored-', '')

457 retval = retval.replace('redacted-', '')

458 # Anchor the end as 3 dashes, a v, and 3 digits, e.g. ---v001

459 match = re.search(r'^(.*)---v\d{3}$', retval)

460 if match:

461 retval = match.group(1)

462 return retval

463

464

465 def es_waiter(client: 'Elasticsearch', cls, **kwargs) -> None:

466 """Wait for ILM Phase & Step to be reached"""

467 try:

468 waiter = cls(client, **kwargs)

469 waiter.wait()

470 except (

471 KeyError,

472 ValueError,

473 TimeoutError,

474 IlmWaitError,

475 NotFoundError,

476 ) as wait_err:

477 msg = f'{cls.__name__}: wait for completion failed: {kwargs}'

478 raise e.BadClientResult(msg, wait_err)