GuardAI/patch_route.py at master · aadyasingh2005/GuardAI · GitHub

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
import os
f = open(os.path.join('backend', 'routes', 'protection.py'), 'r', encoding='utf-8')
content = f.read()
f.close()

old_patterns = [
    "method = options.get('method', 'msap')\r\n        if method == 'blurguard':\r\n            strength = options.get('strength', 'medium')\r\n            result = protection_service.protect_against_qwen_edit(image, strength=strength, options=options)\r\n        else:\r\n            result = protection_service.protect_image(image, options)",
    "method = options.get('method', 'msap')\n        if method == 'blurguard':\n            strength = options.get('strength', 'medium')\n            result = protection_service.protect_against_qwen_edit(image, strength=strength, options=options)\n        else:\n            result = protection_service.protect_image(image, options)",
]

new_block = """method = options.get('method', 'unified')
        strength = options.get('strength', 'medium')
        if method == 'unified':
            result = protection_service.protect_unified(image, strength=strength, options=options)
        elif method == 'blurguard':
            result = protection_service.protect_against_qwen_edit(image, strength=strength, options=options)
        elif method == 'aberration':
            result = protection_service.protect_with_aberration(image, strength=strength, options=options)
        else:
            result = protection_service.protect_image(image, options)"""

patched = False
for old in old_patterns:
    if old in content:
        content = content.replace(old, new_block)
        patched = True
        break

if patched:
    f = open(os.path.join('backend', 'routes', 'protection.py'), 'w', encoding='utf-8')
    f.write(content)
    f.close()
    print('PATCHED OK')
else:
    print('Pattern not found, doing manual line-based approach')
    lines = content.split('\n')
    new_lines = []
    skip = 0
    for i, line in enumerate(lines):
        if skip > 0:
            skip -= 1
            continue
        if "method = options.get('method'" in line.replace('\r','') and 'msap' in line:
            # Count lines to replace
            for j in range(i+1, min(i+10, len(lines))):
                if 'protect_image' in lines[j]:
                    skip = j - i
                    break
            new_lines.append("        method = options.get('method', 'unified')")
            new_lines.append("        strength = options.get('strength', 'medium')")
            new_lines.append("        if method == 'unified':")
            new_lines.append("            result = protection_service.protect_unified(image, strength=strength, options=options)")
            new_lines.append("        elif method == 'blurguard':")
            new_lines.append("            result = protection_service.protect_against_qwen_edit(image, strength=strength, options=options)")
            new_lines.append("        elif method == 'aberration':")
            new_lines.append("            result = protection_service.protect_with_aberration(image, strength=strength, options=options)")
            new_lines.append("        else:")
            new_lines.append("            result = protection_service.protect_image(image, options)")
            patched = True
        else:
            new_lines.append(line)

    if patched:
        f = open(os.path.join('backend', 'routes', 'protection.py'), 'w', encoding='utf-8')
        f.write('\n'.join(new_lines))
        f.close()
        print('PATCHED OK (line-based)')
    else:
        print('FAILED - could not find pattern')