aravagarwal/CodeCloakPII

1 2from copy import deepcopy3 4import gradio as gr5from presidio_anonymizer import OperatorConfig  6from difflib import Differ7from code_diff import difference as code_diff8 9from anoninterface import MAX_LENGTH, MAX_STRINGS_IN_SAMPLE, analyzer, anonymizer10from anoninterface.redaction import generate_constant_id, generate_hash_id11from anoninterface.treeutils import syntax_aware_replace_with_lambda, visualize_tree_sitter_matches, TREE_SITTER_QUERY_NAME_TO_QUERY12from anoninterface.identification import anonymize_string, analyze_string13 14 15 16 17 18 19 20def visualize_redact(source_text):21    queries = ["(comment) @comment", "(string_content) @string_content"]22    #print(text_to_redact)23    for query in queries:24      text_to_redact = syntax_aware_replace_with_lambda(source_text, query, generate_constant_id, anonymize_string)25 26    queries2 = [ "(identifier) @identifier"]27    for query in queries2:28      text_to_redact = syntax_aware_replace_with_lambda(text_to_redact, query, generate_hash_id, anonymize_string)29    30    return text_to_redact31 32 33simple_example = """34# My name is Inigo Montoya35Luke = 136Luke += 137Jane = 138Jane += 239\"\"\"40Monty41Python42Holy43Grail44\"\"\"45real_user_string = \"My name is John Smith\"46 47"""48 49 50 51 52with gr.Blocks() as demo:53    54    gr.Markdown("## Welcome to CodeCloak, an Interactive PII Removal Tool.")55 56    gr.Markdown("We will first pick out the sections of code we want to anonymize using Tree-Sitter, an incremental parser.")57 58    gr.Markdown("Please place the Python code you want to analyze here:")59    with gr.Row():60        inp = gr.Code(value=simple_example, label="Source Text")61    62    gr.Markdown("Please select the elements you want to identify:")63    with gr.Row():64        query_list = gr.Dropdown(choices=list(TREE_SITTER_QUERY_NAME_TO_QUERY.keys()), value=["comment","string","name"],multiselect=True)65    66 67    gr.Markdown("You can either **analyze** text to see what the tool detects as PII, or **anonymize** to get a diff of what the tool changes.")68 69    with gr.Tab("Analyze"):70 71        btn = gr.Button("Locate useful sections")72        73        with gr.Row():74            stageoneoutput = gr.HighlightedText(label="Simple Vis",combine_adjacent=True)75 76 77 78        gr.Markdown("Here are the strings we shall run our NER tool against:")79 80        btn2 = gr.Button("Identify PII")81 82 83        84        85        input_textboxes = []86        output_highlights = []87        for i in range(MAX_STRINGS_IN_SAMPLE):88            with gr.Row(visible=True) as lookatme:89                t = gr.Code(visible=False)90                input_textboxes.append(t)91                output_t = gr.HighlightedText(visible=False)92                output_highlights.append(output_t)93        print('inputs',len(output_highlights))94        print('outputs',len(input_textboxes))95 96        def wrap_visualize_tree_sitter_matches(inputv, query_list):97            labelled_text, capture_list = visualize_tree_sitter_matches(inputv, query_list)98            elems_to_capture = sorted(list(set([elem[0].text.decode('utf8') for elem in capture_list])), key = lambda x: len(x))99            global TOTAL_ITEM_COUNT100            TOTAL_ITEM_COUNT = len(elems_to_capture)101            return [gr.Code(elem, visible=True) for elem in elems_to_capture] + [gr.Code(visible=False)]*(MAX_STRINGS_IN_SAMPLE-len(elems_to_capture))102 103        def wrap_visualize_tree_sitter_matches2(inputv, query_list):104            labelled_text, capture_list = visualize_tree_sitter_matches(inputv, query_list)105            return labelled_text106 107 108 109        110        btn.click(fn=wrap_visualize_tree_sitter_matches2, inputs=[inp, query_list], outputs=[stageoneoutput])111        btn.click(fn=wrap_visualize_tree_sitter_matches, inputs=[inp, query_list], outputs=input_textboxes)112        113        def wrap_model_inference(*input_textboxes):114            num_good_boxes = len([elem for elem in input_textboxes if elem != ''])115            iterative_boxes = [gr.HighlightedText([('No PII',None)], visible=True)]*num_good_boxes + [gr.HighlightedText(visible = False)]*(MAX_STRINGS_IN_SAMPLE-num_good_boxes)116            for elem in range(len(input_textboxes)):117                if len(input_textboxes[elem]) > 0:118                    analysis = analyze_string(input_textboxes[elem])119                    iterative_boxes[elem] = gr.HighlightedText(analysis, visible=True)120                    yield iterative_boxes121                else:122                    iterative_boxes[elem] = gr.HighlightedText(visible=False)123                    yield iterative_boxes124            iterative_boxes125 126        127 128        btn2.click(fn=wrap_model_inference, inputs=input_textboxes, outputs=output_highlights)    129 130    with gr.Tab("Anonymize"):   131        gr.Markdown("Now, we will apply the redactions to these matches, and return the redacted text:")132 133        def wrap_final_output(query_list, inp):134            queries_constant = ["(comment) @comment", "(string_content) @string_content"]135            queries_hash = ["(identifier) @identifier"]136 137            def diff_texts(text1, text2):138                d = Differ()139                return [140                    (token[2:], token[0] if token[0] != " " else None)141                    for token in d.compare(text1, text2)142                ]143            text_to_redact = inp144            for query in query_list:145                query = TREE_SITTER_QUERY_NAME_TO_QUERY[query]146                text_to_redact = syntax_aware_replace_with_lambda(text_to_redact, query, generate_constant_id if query in queries_constant else generate_hash_id, anonymize_string  )147 148            return gr.Code(text_to_redact, visible=True), gr.HighlightedText(diff_texts(inp, text_to_redact),visible=True, combine_adjacent=True), gr.Code(inp, visible=True)149 150        with gr.Row():151            with gr.Column():152                gr.Markdown("Before:")153                initial_text = gr.Code(visible = False)154                gr.Markdown("After:")155                final_text = gr.Code(visible = False)156            final_diff = gr.HighlightedText(visible = False)157 158        btn3 = gr.Button("Anonymize and Compare", visible = True)159        btn3.click(wrap_final_output, inputs=[query_list, inp], outputs = [final_text, final_diff, initial_text])160 161 162 163 164 165 166demo.launch(share=True)167