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

1"""Main app definition"""

2

3# pylint: disable=broad-exception-caught,R0913

4 import typing as t

5 import logging

6 from es_pii_tool.exceptions import FatalError, MissingIndex

7 from es_pii_tool.job import Job

8 from es_pii_tool.redacters.index import RedactIndex

9 from es_pii_tool.task import Task

10 from es_pii_tool.helpers.elastic_api import get_hits

11 from es_pii_tool.helpers.utils import end_it, get_redactions

12

13 if t.TYPE_CHECKING:

14 from elasticsearch8 import Elasticsearch

15

16 logger = logging.getLogger(__name__)

17

18

19 class PiiTool:

20 """Elasticsearch PII Tool"""

21

22 def __init__(

23 self,

24 client: 'Elasticsearch',

25 tracking_index: str,

26 redaction_file: str = '',

27 redaction_dict: t.Union[t.Dict, None] = None,

28 dry_run: bool = False,

29 ):

30 if redaction_dict is None:

31 redaction_dict = {}

32 logger.debug('Redactions file: %s', redaction_file)

33 self.counter = 0

34 self.client = client

35 self.redactions = get_redactions(redaction_file, redaction_dict)

36 self.tracking_index = tracking_index

37 self.dry_run = dry_run

38

39 def verify_doc_count(self, job: Job) -> bool:

40 """Verify that expected_docs and the hits from the query have the same value

41

42 :param job: The job object for the present redaction run

43

44 :type job: :py:class:`~.app.tracking.Job`

45

46 :rtype: None

47 :returns: No return value

48 """

49 task = Task(job, task_id=f'PRE---{job.name}---DOC-COUNT-VERIFICATION')

50 success = False

51 errors = False

52 if task.finished():

53 return True # We're done already

54 # Log task start

55 task.begin()

56 hits = get_hits(self.client, job.config['pattern'], job.config['query'])

57 msg = f'{hits} hit(s)'

58 logger.debug(msg)

59 task.add_log(msg)

60 logger.info("Checking expected document count...")

61 zeromsg = (

62 f"For index pattern {job.config['pattern']}, with query "

63 f"{job.config['query']} 'expected_docs' is {job.config['expected_docs']} "

64 f"but query results is {hits} matches."

65 )

66 if job.config['expected_docs'] == hits:

67 msg = (

68 f'Query result hits: {hits} matches expected_docs: '

69 f'{job.config["expected_docs"]}'

70 )

71 logger.debug(msg)

72 task.add_log(msg)

73 success = True

74 if hits == 0:

75 logger.critical(zeromsg)

76 logger.info('Continuing to next configuration block (if any)')

77 success = False

78 else:

79 logger.critical(zeromsg)

80 logger.info('Continuing to next configuration block (if any)')

81 if not success:

82 errors = True

83 task.add_log(zeromsg)

84 task.end(success, errors=errors)

85 return success

86

87 def iterate_indices(self, job: Job) -> bool:

88 """Iterate over every index in job.indices"""

89 all_succeeded = True

90 for idx in job.indices:

91 task = Task(job, index=idx, id_suffix='PARENT-TASK')

92 # First check to see if idx has been touched as part of a previous run

93 if task.finished():

94 continue # This index has already been verified

95 task.begin()

96 task_success = False

97 try:

98 msg = f'Iterating per index: Index {idx} of {job.indices}'

99 logger.debug(msg)

100 task.add_log(msg)

101 redact = RedactIndex(idx, job, self.counter)

102 redact.run()

103 task_success = redact.success

104 self.counter = redact.counter

105 logger.debug('RESULT: %s', task_success)

106 except MissingIndex as err:

107 logger.critical(err)

108 raise FatalError(f'Index {err.missing} not found.', err) from err

109 except FatalError as err:

110 logger.critical('Fatal upstream error encountered: %s', err.message)

111 raise FatalError('We suffered a fatal upstream error', err) from err

112 end_it(task, task_success)

113 if not task.completed:

114 all_succeeded = False

115 job.add_log(f'Unable to complete task {task.task_id}')

116 return all_succeeded

117

118 def iterate_configuration(self) -> None:

119 """Iterate over every configuration block in self.redactions"""

120 logger.debug('Full redactions object from config: %s', self.redactions)

121 for config_block in self.redactions['redactions']: # type: ignore

122 job_success = True

123 # Reset counter to zero for each full iteration

124 self.counter = 0

125 if self.dry_run:

126 logger.info("DRY-RUN MODE ENABLED. No data will be changed.")

127

128 # There's really only 1 root-level key for each configuration block,

129 # and that's job_id

130 job_name = list(config_block.keys())[0]

131 args = (self.client, self.tracking_index, job_name, config_block[job_name])

132 job = Job(*args, dry_run=self.dry_run)

133 if job.finished():

134 continue

135 job.begin()

136 if not self.verify_doc_count(job):

137 # This configuration block can't go further because of the mismatch

138 job_success = False

139 end_it(job, job_success)

140 continue

141

142 job_success = self.iterate_indices(job)

143 # At this point, self.counter should be equal to total, indicating that we

144 # matched expected_docs. We should therefore register that the job was

145 # successful, if we have reached this point with no other errors having

146 # interrupted the process.

147

148 end_it(job, job_success)

149

150 def run(self) -> None:

151 """Do the thing"""

152 logger.info('PII scrub initiated')

153 self.iterate_configuration()