import re
import os
from pathlib import Path
def update_file(file_path):
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
original_content = content
updated = False
if 'use gun::' in content or 'use gun::' in content:
if 'use chia_bls::' not in content:
use_pattern = r'(use [^;]+;)'
matches = list(re.finditer(use_pattern, content))
if matches:
last_use = matches[-1]
insert_pos = last_use.end()
content = content[:insert_pos] + '\nuse chia_bls::{SecretKey, PublicKey};' + content[insert_pos:]
updated = True
if 'Gun::new()' in content:
pattern = r'(\s+)(let\s+gun\s*=\s*)Gun::new\(\);'
replacement = r'\1// Generate BLS key pair\n\1let secret_key = SecretKey::from_seed(&[0u8; 32]);\n\1let public_key = secret_key.public_key();\n\1\2Gun::new(secret_key, public_key);'
new_content = re.sub(pattern, replacement, content)
if new_content != content:
content = new_content
updated = True
pattern2 = r'(\s+)(let\s+\w+\s*=\s*Arc::new\()Gun::new\(\)\);'
replacement2 = r'\1// Generate BLS key pair\n\1let secret_key = SecretKey::from_seed(&[0u8; 32]);\n\1let public_key = secret_key.public_key();\n\1\2Gun::new(secret_key, public_key));'
new_content = re.sub(pattern2, replacement2, content)
if new_content != content:
content = new_content
updated = True
if 'Gun::with_options(' in content:
pattern1 = r'(\s+)(let\s+secret_key\d*\s*=\s*SecretKey::from_seed\([^)]+\);\s*\n\s*let\s+public_key\d*\s*=\s*secret_key\d*\.public_key\(\);)?\s*(let\s+\w+\s*=\s*match\s+)?Gun::with_options\((\w+)\)\.await'
lines = content.split('\n')
new_lines = []
i = 0
key_counter = 1
while i < len(lines):
line = lines[i]
if 'Gun::with_options(' in line and 'SecretKey::from_seed' not in '\n'.join(lines[max(0, i-5):i]):
indent = len(line) - len(line.lstrip())
indent_str = ' ' * indent
match = re.search(r'Gun::with_options\((\w+)\)', line)
if match:
options_var = match.group(1)
key_gen = f'{indent_str}let secret_key{key_counter} = SecretKey::from_seed(&[{key_counter}u8; 32]);\n'
key_gen += f'{indent_str}let public_key{key_counter} = secret_key{key_counter}.public_key();\n'
new_lines.append(key_gen.rstrip())
line = line.replace(
f'Gun::with_options({options_var})',
f'Gun::with_options(secret_key{key_counter}, public_key{key_counter}, {options_var})'
)
key_counter += 1
updated = True
new_lines.append(line)
i += 1
if updated:
content = '\n'.join(new_lines)
if updated and content != original_content:
with open(file_path, 'w', encoding='utf-8') as f:
f.write(content)
return True
return False
def main():
base_dir = Path('.')
files_to_update = []
for pattern in ['examples/**/*.rs', 'tests/**/*.rs']:
files_to_update.extend(base_dir.glob(pattern))
updated_count = 0
for file_path in sorted(files_to_update):
if update_file(str(file_path)):
print(f"Updated: {file_path}")
updated_count += 1
print(f"\nTotal files updated: {updated_count}")
if __name__ == '__main__':
main()