File: sanitizers.py

package info (click to toggle)
python-azure 20250603%2Bgit-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 851,724 kB
  • sloc: python: 7,362,925; ansic: 804; javascript: 287; makefile: 195; sh: 145; xml: 109
file content (829 lines) | stat: -rw-r--r-- 40,778 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
from enum import Enum
import json
from typing import Dict, List, Optional

from .config import PROXY_URL
from .helpers import get_http_client, get_recording_id, is_live, is_live_and_not_recording


class Sanitizer(str, Enum):
    """Sanitizers that can be applied to recordings."""

    BODY_KEY = "BodyKeySanitizer"
    BODY_REGEX = "BodyRegexSanitizer"
    BODY_STRING = "BodyStringSanitizer"
    CONTINUATION = "ContinuationSanitizer"
    GENERAL_REGEX = "GeneralRegexSanitizer"
    GENERAL_STRING = "GeneralStringSanitizer"
    HEADER_REGEX = "HeaderRegexSanitizer"
    HEADER_STRING = "HeaderStringSanitizer"
    OAUTH_RESPONSE = "OAuthResponseSanitizer"
    REMOVE_HEADER = "RemoveHeaderSanitizer"
    URI_REGEX = "UriRegexSanitizer"
    URI_STRING = "UriStringSanitizer"
    URI_SUBSCRIPTION_ID = "UriSubscriptionIdSanitizer"


# This file contains methods for adjusting many aspects of test proxy behavior:
#
# - Sanitizers: record stand-in values to hide secrets and/or enable playback when behavior is inconsistent
# - Transforms: extend test proxy functionality by changing how recordings are processed in playback mode
# - Matchers: modify the conditions that are used to match request and response content with recorded values
# - Recording options: further customization for advanced scenarios, such as providing certificates to the transport
#
# Methods for a given category are grouped together under a header containing more details.


def set_default_function_settings() -> None:
    """Resets sanitizers, matchers, and transforms for the test proxy to their default settings, for the current test.

    This will reset any setting customizations for a single test. This must be called during test case execution, rather
    than at a session, module, or class level. To reset setting customizations for all tests, use
    `set_default_session_settings` instead.
    """

    x_recording_id = get_recording_id()
    if x_recording_id is None:
        raise RuntimeError(
            "This method must be called during test case execution. To reset test proxy settings at a session level, "
            "use `set_default_session_settings` instead."
        )
    _send_reset_request({"x-recording-id": x_recording_id})


def set_default_session_settings() -> None:
    """Resets sanitizers, matchers, and transforms for the test proxy to their default settings, for all tests.

    This will reset any setting customizations for an entire test session. To reset setting customizations for a single
    test -- which is recommended -- use `set_default_function_settings` instead.
    """

    _send_reset_request({})


# ----------MATCHERS----------
#
# A matcher is applied during a playback session. The default matcher matches a request on headers, URI, and the body.
#
# This is the least used customization as most adjustments to matching really come down to sanitizing properly before
# storing the recording. Further, when using this customization, it is recommended that one registers matchers during
# individual test case execution so that the adjusting matching only occurs for a specific recording during playback.
#
# ----------------------------


def set_bodiless_matcher() -> None:
    """Adjusts the "match" operation to EXCLUDE the body when matching a request to a recording's entries.

    This method should be called during test case execution, rather than at a session, module, or class level.
    """

    x_recording_id = get_recording_id()
    _send_matcher_request("BodilessMatcher", {"x-recording-id": x_recording_id})


def set_custom_default_matcher(**kwargs) -> None:
    """Exposes the default matcher in a customizable way.

    All optional settings are safely defaulted. This means that providing zero additional configuration will produce a
    sanitizer that is functionally identical to the default.

    :keyword bool compare_bodies: True to enable body matching (default behavior), or False to disable body matching.
    :keyword str excluded_headers: A comma separated list of headers that should be excluded during matching. The
        presence of these headers will not be taken into account while matching. Should look like
        "Authorization, Content-Length", for example.
    :keyword str ignored_headers: A comma separated list of headers that should be ignored during matching. The header
        values won't be matched, but the presence of these headers will be taken into account while matching. Should
        look like "Authorization, Content-Length", for example.
    :keyword bool ignore_query_ordering: By default, the test proxy does not sort query params before matching. Setting
        to True will sort query params alphabetically before comparing URIs.
    :keyword str ignored_query_parameters: A comma separated list of query parameters that should be ignored during
        matching. The parameter values won't be matched, but the presence of these parameters will be taken into account
        while matching. Should look like "param1, param2", for example.
    """

    x_recording_id = get_recording_id()
    request_args = _get_request_args(**kwargs)
    _send_matcher_request("CustomDefaultMatcher", {"x-recording-id": x_recording_id}, request_args)


def set_headerless_matcher() -> None:
    """Adjusts the "match" operation to ignore header differences when matching a request.

    Be aware that wholly ignoring headers during matching might incur unexpected issues down the line. This method
    should be called during test case execution, rather than at a session, module, or class level.
    """

    x_recording_id = get_recording_id()
    _send_matcher_request("HeaderlessMatcher", {"x-recording-id": x_recording_id})


# ----------SANITIZERS----------
#
# A sanitizer is applied to recordings in two locations:
#
# - Before they are saved. (Affects the session as a whole as well as individual entries)
# - During playback, when a request comes in from a client. This means that only individual entry sanitizers apply in
#   this case.
#
# Sanitizers are optionally prefixed with a title that indicates where each sanitizer applies. These prefixes are:
#
# - Uri
# - Header
# - Body
#
# For example, A sanitizer prefixed with Body will only ever operate on the request/response body. The target URI and
# request/response headers will be left unaffected.
#
# ------------------------------


def add_body_key_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that offers regex update of a specific JTokenPath within a returned body.

    For example, "TableName" within a json response body having its value replaced by whatever substitution is offered.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str json_path: The SelectToken path (which could possibly match multiple entries) that will be used to
        select JTokens for value replacement.
    :keyword str value: The substitution value.
    :keyword str regex: A regex. Can be defined as a simple regex replace OR if groupForReplace is set, a substitution
        operation. Defaults to replacing the entire string.
    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
        a simple replacement operation.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("BodyKeySanitizer", request_args, {"x-recording-id": x_recording_id})


def add_body_regex_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that offers regex replace within a returned body.

    Specifically, this means regex applying to the raw JSON. If you are attempting to simply replace a specific key, the
    BodyKeySanitizer is probably the way to go.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str value: The substitution value.
    :keyword str regex: A regex. Can be defined as a simple regex, or if a ``group_for_replace`` is provided, a
        substitution operation.
    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
        a simple replacement operation.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("BodyRegexSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_body_string_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that cleans request and response bodies via straightforward string replacement.

    Specifically, this replacement applies to the raw JSON of a body. If you are attempting to simply replace a specific
    key, add_body_key_sanitizer is probably more suitable.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str value: The substitution value.
    :keyword str target: A target string. This could contain special regex characters like "?()+*" but they will be
        treated as a literal.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("BodyStringSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_continuation_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that's used to anonymize private keys in response/request pairs.

    For instance, a request hands back a "sessionId" that needs to be present in the next request. Supports "all further
    requests get this key" as well as "single response/request pair". Defaults to maintaining same key for rest of
    recording.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str key: The name of the header whose value will be replaced from response -> next request.
    :keyword str method: The method by which the value of the targeted key will be replaced. Defaults to guid
        replacement.
    :keyword str reset_after_first: Do we need multiple pairs replaced? Or do we want to replace each value with the
        same value?
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("ContinuationSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_general_regex_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that offers a general regex replace across request/response Body, Headers, and URI.

    For the body, this means regex applying to the raw JSON.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str value: The substitution value.
    :keyword str regex: A regex. Can be defined as a simple regex, or if a ``group_for_replace`` is provided, a
        substitution operation.
    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
        a simple replacement operation.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("GeneralRegexSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_general_string_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that cleans request and response URIs, headers, and bodies via string replacement.

    This sanitizer offers a value replace across request/response bodies, headers, and URIs. For the body, this means a
    string replacement applied directly to the raw JSON.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str value: The substitution value.
    :keyword str target: A target string. This could contain special regex characters like "?()+*" but they will be
        treated as a literal.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("GeneralStringSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_header_regex_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that offers regex replace on returned headers.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    Can be used for multiple purposes: 1) To replace a key with a specific value, do not set "regex" value. 2) To do a
    simple regex replace operation, define arguments "key", "value", and "regex". 3) To do a targeted substitution of a
    specific group, define all arguments "key", "value", and "regex".

    :keyword str key: The name of the header we're operating against.
    :keyword str value: The substitution or whole new header value, depending on "regex" setting.
    :keyword str regex: A regex. Can be defined as a simple regex, or if a ``group_for_replace`` is provided, a
        substitution operation.
    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
        a simple replacement operation.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("HeaderRegexSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_header_string_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that cleans headers in a recording via straightforward string replacement.

    This sanitizer ONLY applies to the request/response headers -- bodies and URIs are left untouched.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str key: The name of the header we're operating against.
    :keyword str target: A target string. This could contain special regex characters like "?()+*" but they will be
        treated as a literal.
    :keyword str value: The substitution value.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("HeaderStringSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_oauth_response_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that cleans out all request/response pairs that match an oauth regex in their URI.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    _send_sanitizer_request("OAuthResponseSanitizer", {}, {"x-recording-id": x_recording_id})


def add_remove_header_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that removes specified headers before saving a recording.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str headers: A comma separated list. Should look like "Location, Transfer-Encoding" or something along
        those lines. Don't worry about whitespace between the commas separating each key. They will be ignored.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("RemoveHeaderSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_uri_regex_sanitizer(**kwargs) -> None:
    """Registers a sanitizer for cleaning URIs via regex.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str value: The substitution value.
    :keyword str regex: A regex. Can be defined as a simple regex, or if a ``group_for_replace`` is provided, a
        substitution operation.
    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
        a simple replacement operation.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("UriRegexSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_uri_string_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that cleans URIs via straightforward string replacement.

    Runs a simple string replacement against the request/response URIs.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str value: The substitution value.
    :keyword str target: A target string. This could contain special regex characters like "?()+*" but they will be
        treated as a literal.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("UriStringSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_uri_subscription_id_sanitizer(**kwargs) -> None:
    """Registers a sanitizer that replaces subscription IDs in URIs.

    This sanitizer ONLY affects the URI of a request/response pair. Subscription IDs are replaced with
    "00000000-0000-0000-0000-000000000000" by default.

    :keyword bool function_scoped: Whether or not the sanitizer should apply only to the test function being currently
        executed. If set to True, the sanitizer must be called during a specific test's execution. Defaults to False.
    :keyword str value: The fake subscription ID that will be placed where the real one is in the real request.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    function_scoped = kwargs.pop("function_scoped", False)
    x_recording_id = get_recording_id() if function_scoped else None
    request_args = _get_request_args(**kwargs)
    _send_sanitizer_request("UriSubscriptionIdSanitizer", request_args, {"x-recording-id": x_recording_id})


def add_batch_sanitizers(sanitizers: Dict[str, List[Optional[Dict[str, str]]]], headers: Optional[Dict] = None) -> None:
    """Registers a batch of sanitizers at once.

    If live tests are being run with recording turned off via the AZURE_SKIP_LIVE_RECORDING environment variable, no
    request will be sent.

    :param sanitizers: A group of sanitizers to add, as a dictionary. Keys should be sanitizer names (from the Sanitizer
        enum) and values should be lists containing dictionaries of sanitizer constructor parameters. The parameters
        should be formatted as key-value pairs aligning with keyword-only arguments to sanitizer methods.
    :type sanitizers: dict[str, list[Optional[dict]]]
    """

    if is_live_and_not_recording():
        return

    data = []  # Body content to populate with multiple sanitizer definitions

    for sanitizer in sanitizers:
        # Iterate over each instance of the particular sanitizer (e.g. each body regex sanitizer)
        for sanitizer_instance in sanitizers[sanitizer]:
            sanitizer_definition = {"Name": sanitizer}
            if sanitizer_instance:
                sanitizer_definition.update({"Body": _get_request_args(**sanitizer_instance)})
            data.append(sanitizer_definition)

    headers_to_send = {"Content-Type": "application/json"}
    x_recording_id = get_recording_id()
    if x_recording_id:
        headers_to_send["x-recording-id"] = x_recording_id
    if headers is not None:
        for key in headers:
            if headers[key] is not None:
                headers_to_send[key] = headers[key]

    http_client = get_http_client()
    http_client.request(
        method="POST",
        url="{}/Admin/AddSanitizers".format(PROXY_URL),
        headers=headers_to_send,
        body=json.dumps(data).encode("utf-8"),
    )


def remove_batch_sanitizers(sanitizers: List[str], headers: Optional[Dict] = None) -> None:
    """Removes a batch of sanitizers.

    Sanitizers are denoted by their ID, which is a string. This method will remove all sanitizers with the provided
    IDs.

    :param sanitizers: A list of sanitizer IDs to remove.
    :type sanitizers: list[str]
    :param headers: Optional headers to include in the request.
    :type headers: dict
    """

    if is_live_and_not_recording():
        return

    data = {"Sanitizers": sanitizers}

    headers_to_send = {"Content-Type": "application/json"}
    if headers is not None:
        for key in headers:
            if headers[key] is not None:
                headers_to_send[key] = headers[key]

    http_client = get_http_client()
    http_client.request(
        method="POST",
        url="{}/Admin/RemoveSanitizers".format(PROXY_URL),
        headers=headers_to_send,
        body=json.dumps(data).encode("utf-8"),
    )


# ----------TRANSFORMS----------
#
# A transform extends functionality of the test proxy by applying to responses just before they are returned during
# playback mode.
#
# ------------------------------


def add_api_version_transform() -> None:
    """Registers a transform that copies a request's "api-version" header onto the response before returning it.

    This method should be called during test case execution, rather than at a session, module, or class level.
    """

    x_recording_id = get_recording_id()
    _send_transform_request("ApiVersionTransform", {}, {"x-recording-id": x_recording_id})


def add_client_id_transform() -> None:
    """Registers a transform that copies a request's "x-ms-client-id" header onto the response before returning it.

    This method should be called during test case execution, rather than at a session, module, or class level.
    """

    x_recording_id = get_recording_id()
    _send_transform_request("ClientIdTransform", {}, {"x-recording-id": x_recording_id})


def add_header_transform(**kwargs) -> None:
    """Registers a transform that sets a header in a response.

    This method should be called during test case execution, rather than at a session, module, or class level.

    :keyword str key: The key for the header.
    :keyword str value: The value for the header.
    :keyword str condition: A condition that dictates when this sanitizer applies to a request/response pair. The
        content of this should be a JSON object that contains configuration keys. Currently, that only includes the key
        "uriRegex". This translates to an object that looks like '{ "uriRegex": "when this regex matches, apply the
        sanitizer" }'. Defaults to "apply always".
    """

    x_recording_id = get_recording_id()
    request_args = _get_request_args(**kwargs)
    _send_transform_request("HeaderTransform", request_args, {"x-recording-id": x_recording_id})


def add_storage_request_id_transform() -> None:
    """Registers a transform that ensures a response's "x-ms-client-request-id" header matches the request's.

    This method should be called during test case execution, rather than at a session, module, or class level.
    """

    x_recording_id = get_recording_id()
    _send_transform_request("StorageRequestIdTransform", {}, {"x-recording-id": x_recording_id})


# ----------RECORDING OPTIONS----------
#
# Recording options enable customization beyond what is offered by sanitizers, matchers, and transforms. These are
# intended for advanced scenarios and are generally not applicable.
#
# -------------------------------------


def set_function_recording_options(**kwargs) -> None:
    """Sets custom recording options for the current test only.

    This must be called during test case execution, rather than at a session, module, or class level. To set recording
    options for all tests, use `set_session_recording_options` instead.

    :keyword bool handle_redirects: The test proxy performs transparent follow directs by default. That means
        that if the initial request sent through the test proxy results in a 3XX redirect status, the test proxy will
        follow. Setting `handle_redirects` to False will instead make the test proxy return that redirect response to
        the client and allow it to handle the redirect.
    :keyword str context_directory: This changes the "root" path that the test proxy uses when loading a recording.
    :keyword certificates: A list of `PemCertificate`s. Any number of certificates is allowed.
    :type certificates: Iterable[PemCertificate]
    :keyword str tls_certificate: The public key portion of a TLS certificate, as a string. This is used specifically so
        that an SSL connection presenting a non-standard certificate can still be validated.
    """

    x_recording_id = get_recording_id()
    request_args = _get_recording_option_args(**kwargs)
    _send_recording_options_request(request_args, {"x-recording-id": x_recording_id})


def set_session_recording_options(**kwargs) -> None:
    """Sets custom recording options for all tests.

    This will set the specified recording options for an entire test session. To set recording options for a single test
    -- which is recommended -- use `set_function_recording_options` instead.

    :keyword bool handle_redirects: The test proxy performs transparent follow directs by default. That means
        that if the initial request sent through the test proxy results in a 3XX redirect status, the test proxy will
        follow. Setting `handle_redirects` to False will instead make the test proxy return that redirect response to
        the client and allow it to handle the redirect.
    :keyword str context_directory: This changes the "root" path that the test proxy uses when loading a recording.
    :keyword certificates: A list of `PemCertificate`s. Any number of certificates is allowed.
    :type certificates: Iterable[PemCertificate]
    :keyword str tls_certificate: The public key portion of a TLS certificate, as a string. This is used specifically so
        that an SSL connection presenting a non-standard certificate can still be validated.
    """

    request_args = _get_recording_option_args(**kwargs)
    _send_recording_options_request(request_args)


class PemCertificate:
    """Represents a PEM certificate that can be sent to and used by the test proxy.

    :param str data: The content of the certificate, as a string.
    :param str key: The certificate key, as a string.
    """

    def __init__(self, data: str, key: str) -> None:
        self.data = data
        self.key = key


# ----------HELPERS----------


def _get_recording_option_args(**kwargs) -> Dict:
    """Returns a dictionary of recording option request arguments, formatted for test proxy consumption."""

    certificates = kwargs.pop("certificates", None)
    tls_certificate = kwargs.pop("tls_certificate", None)
    tls_certificate_host = kwargs.pop("tls_certificate_host", None)
    request_args = _get_request_args(**kwargs)

    if certificates or tls_certificate:
        transport = {}

        if certificates:
            cert_pairs = [{"PemValue": cert.data, "PemKey": cert.key} for cert in certificates]
            transport["Certificates"] = cert_pairs

        if tls_certificate:
            transport["TLSValidationCert"] = tls_certificate

        if tls_certificate_host:
            transport["TSLValidationCertHost"] = tls_certificate_host

        request_args["Transport"] = transport

    return request_args


def _get_request_args(**kwargs) -> Dict:
    """Returns a dictionary of request arguments, formatted for test proxy consumption."""

    request_args = {}
    if "compare_bodies" in kwargs:
        request_args["compareBodies"] = kwargs.get("compare_bodies")
    if "condition" in kwargs:
        request_args["condition"] = kwargs.get("condition")
    if "context_directory" in kwargs:
        request_args["ContextDirectory"] = kwargs.get("context_directory")
    if "excluded_headers" in kwargs:
        request_args["excludedHeaders"] = kwargs.get("excluded_headers")
    if "group_for_replace" in kwargs:
        request_args["groupForReplace"] = kwargs.get("group_for_replace")
    if "handle_redirects" in kwargs:
        request_args["HandleRedirects"] = kwargs.get("handle_redirects")
    if "headers" in kwargs:
        request_args["headersForRemoval"] = kwargs.get("headers")
    if "ignored_headers" in kwargs:
        request_args["ignoredHeaders"] = kwargs.get("ignored_headers")
    if "ignore_query_ordering" in kwargs:
        request_args["ignoreQueryOrdering"] = kwargs.get("ignore_query_ordering")
    if "ignored_query_parameters" in kwargs:
        request_args["ignoredQueryParameters"] = kwargs.get("ignored_query_parameters")
    if "json_path" in kwargs:
        request_args["jsonPath"] = kwargs.get("json_path")
    if "key" in kwargs:
        request_args["key"] = kwargs.get("key")
    if "method" in kwargs:
        request_args["method"] = kwargs.get("method")
    if "regex" in kwargs:
        request_args["regex"] = kwargs.get("regex")
    if "reset_after_first" in kwargs:
        request_args["resetAfterFirst"] = kwargs.get("reset_after_first")
    if "target" in kwargs:
        request_args["target"] = kwargs.get("target")
    if "value" in kwargs:
        request_args["value"] = kwargs.get("value")
    return request_args


def _send_matcher_request(matcher: str, headers: Dict, parameters: Optional[Dict] = None) -> None:
    """Sends a POST request to the test proxy endpoint to register the specified matcher.

    If live tests are being run, no request will be sent.

    :param str matcher: The name of the matcher to set.
    :param dict headers: Any matcher headers, as a dictionary.
    :param parameters: Any matcher constructor parameters, as a dictionary. Defaults to None.
    :type parameters: Optional[dict]
    """

    if is_live():
        return

    headers_to_send = {"x-abstraction-identifier": matcher}
    for key in headers:
        if headers[key] is not None:
            headers_to_send[key] = headers[key]

    http_client = get_http_client()
    http_client.request(
        method="POST",
        url=f"{PROXY_URL}/Admin/SetMatcher",
        headers=headers_to_send,
        body=json.dumps(parameters).encode("utf-8"),
    )


def _send_recording_options_request(parameters: Dict, headers: Optional[Dict] = None) -> None:
    """Sends a POST request to the test proxy endpoint to set the specified recording options.

    If live tests are being run with recording turned off via the AZURE_SKIP_LIVE_RECORDING environment variable, no
    request will be sent.

    :param dict parameters: The recording options, as a dictionary.
    :param headers: Any recording option request headers, as a dictionary. Defaults to None.
    :type headers: Optional[dict]
    """

    if is_live_and_not_recording():
        return

    headers_to_send = {"Content-Type": "application/json"}
    if headers:
        for key in headers:
            if headers[key] is not None:
                headers_to_send[key] = headers[key]

    http_client = get_http_client()
    http_client.request(
        method="POST",
        url=f"{PROXY_URL}/Admin/SetRecordingOptions",
        headers=headers_to_send,
        body=json.dumps(parameters).encode("utf-8"),
    )


def _send_reset_request(headers: Dict) -> None:
    """Sends a POST request to the test proxy endpoint to reset setting customizations.

    If live tests are being run with recording turned off via the AZURE_SKIP_LIVE_RECORDING environment variable, no
    request will be sent.

    :param dict headers: Any reset request headers, as a dictionary.
    """

    if is_live_and_not_recording():
        return

    headers_to_send = {}
    for key in headers:
        if headers[key] is not None:
            headers_to_send[key] = headers[key]

    http_client = get_http_client()
    http_client.request(method="POST", url=f"{PROXY_URL}/Admin/Reset", headers=headers_to_send)


def _send_sanitizer_request(sanitizer: str, parameters: Dict, headers: Optional[Dict] = None) -> None:
    """Sends a POST request to the test proxy endpoint to register the specified sanitizer.

    If live tests are being run with recording turned off via the AZURE_SKIP_LIVE_RECORDING environment variable, no
    request will be sent.

    :param str sanitizer: The name of the sanitizer to add.
    :param dict parameters: The sanitizer constructor parameters, as a dictionary.
    """

    if is_live_and_not_recording():
        return

    headers_to_send = {"x-abstraction-identifier": sanitizer, "Content-Type": "application/json"}
    if headers:
        for key in headers:
            if headers[key] is not None:
                headers_to_send[key] = headers[key]

    http_client = get_http_client()
    http_client.request(
        method="POST",
        url="{}/Admin/AddSanitizer".format(PROXY_URL),
        headers=headers_to_send,
        body=json.dumps(parameters).encode("utf-8"),
    )


def _send_transform_request(transform: str, parameters: Dict, headers: Optional[Dict] = None) -> None:
    """Sends a POST request to the test proxy endpoint to register the specified transform.

    If live tests are being run, no request will be sent.

    :param str transform: The name of the transform to add.
    :param dict parameters: The transform constructor parameters, as a dictionary.
    """

    if is_live():
        return

    headers_to_send = {"x-abstraction-identifier": transform, "Content-Type": "application/json"}
    if headers:
        for key in headers:
            if headers[key] is not None:
                headers_to_send[key] = headers[key]

    http_client = get_http_client()
    http_client.request(
        method="POST",
        url=f"{PROXY_URL}/Admin/AddTransform",
        headers=headers_to_send,
        body=json.dumps(parameters).encode("utf-8"),
    )