convert ida decompiler output into something a bit closer to proper objc syntax

dec.py

code = """your IDA decompiler output here"""

lines = code.split('\n')

import re
from enum import Enum

mfp1 = re.compile('(\S*?) \S*? ([-+])\[(\S*) ([^\]]*)]\(([^\)]*)\)')
depth1_objc_call = re.compile('([-+])\[(\S*) ([^\]]*)]\(([^\)]*)\)')

objc_msg = re.compile('objc_msgSend\(([^,]*), ([^,]*), ([^)]*)\);')
cfstr_to = re.compile('(.*)CFSTR\((.*?\))\)(.*)')
cfstr_frm = re.compile('(.*)#CFSTR#(".*")(.*)')

depth = 0


class VariableType(Enum):
    CLASS = 0
    MSG_REPLY = 1
    MSG_AUTORET = 2
    NORMAL = 64


def proc_objc_msg(statement):
    statement = statement.strip()
    out = statement
    statement = cfstr_to.sub(r'\1#CFSTR#\2\3', statement)
    if ',' in statement.split(', sel_')[1]:
        stat = objc_msg.sub(r'\1 | \2 | \3', statement).split(' | ')
        try:
            recv = stat[0]
            selector = stat[1][4:].replace('_', ':')
            args = stat[2].split(', ')
        except IndexError:
            return out
        
        segments = []
        
        for i, item in enumerate(selector.split(':')):
            if item == "":
                continue
            try:
                argname = args[i]
                segments.append(item + ':' + argname + ' ')
            except IndexError:
                segments.append(item)

        sig = ''.join(segments)
        recv = recv.split(' = ')[-1] # WHERE IS THIS FUCKING EQUAL SIGN COMOING FROM ??!?!??! HOW DOES IT GET HERE? ???!?!?!??!?!??
        out = f'[{recv.strip()} {sig.strip()}];'
    else:
        args = statement.split('Send(')[1].split(', ')
        out = f'[{args[0]} {args[1][:-2].replace("sel_", "")}];'
    out = cfstr_frm.sub(r'\1 CFSTR(\2) \3', out)
    return out



        
def process_objc_decl(objc_decl, depth):
    
    out = ""
    
    words = objc_decl.split(" ")
    if depth == 0:
        argset = mfp1.sub(r'\2(\1) | \3 | \4 | \5', objc_decl).split(' | ')
        decl = argset[0]
        classname = argset[1]
        selector = argset[2]
        args = argset[3].split(', ')
        
        segments = []
        for i, item in enumerate(selector.split(':')):
            if item == "":
                continue
            try:
                argtype = args[i + 2]
                if ' ' in args[i + 2]:
                    argname = args[i + 2].split(' ')[1]
                    argtype = argtype.split(" ")[0]
                else:
                    argname = f'arg{i}'
                segments.append(item + ':' + '(' + argtype + ')' + argname + ' ')
            except IndexError:
                segments.append(item)

        sig = ''.join(segments)
        out = f'{decl} {sig}'
    else:
        if ':' not in objc_decl:
            return objc_decl.strip().split('(', 1)[0][1:]
        argset = depth1_objc_call.sub(r'\2 | \3 | \4', objc_decl[:-1].strip()).split(' | ')
        try:
            classname = argset[0]
            selector = argset[1]
            args = argset[2].split(', ')
            classname = args[0]
        except IndexError:
            return objc_decl
        
        segments = []
        for i, item in enumerate(selector.split(':')):
            if item == "":
                continue
            try:
                argname = args[i+2]
                segments.append(item + ':' + argname + ' ')
            except IndexError:
                segments.append(item)


        sig = ''.join(segments)
        out = f'[{classname} {sig}];'
    return out



variables = {}

variable_declarations = {}

class Variable:
    def __init__(self, statement):
        statement = str(statement)
        self.name = statement.split(' = ', 1)[0]
        self.oname = self.name + ''
        if self.name not in variables:
            if self.name in variable_declarations:
                self.name = variable_declarations[self.name]
        value = statement.split(' = ', 1)[1]
        self.type = VariableType.NORMAL
        self.value = value
        
        if value.strip().startswith('objc_msgSend('):
            self.type = VariableType.MSG_REPLY
            self.value = proc_objc_msg(value)
            
        elif value.strip().startswith('-['):
            self.type = VariableType.MSG_REPLY
            self.value = process_objc_decl(value, 1)
            
        elif value.strip().startswith('objc_retainAutoreleasedReturnValue'):
            self.type = VariableType.MSG_AUTORET
            self.repr = value.strip().split('(')[1][:-2]
            
        
        

processed_lines = []


def add_line(line, depth, indent=4):
    processed_lines.append(f'{" " * (indent * depth)}{line}')
    
    
waiting_variable_msg = None
waited_once = False

in_declaration_section = True


for line in lines:
    line = line.strip().replace(' ;', ';')
    
    if in_declaration_section and depth != 0:
        if '; // ' in line:
            decl = line.strip().split(';')[0]
            name = decl.split(' ')[-1].replace('*', '')
            variable_declarations[name] = decl
            continue
        else:
            in_declaration_section = False
    
    if line.startswith('{'):
        add_line('{', depth)
        depth += 1
    if line.startswith('}'):
        depth -= 1
        add_line('}', depth)
        
        
    if '-[' in line  and '=' not in line or '+[' in line and '=' not in line:
        add_line(process_objc_decl(line, depth), depth)
        
        
    elif ' = ' in line:
        statement = line.strip()
        var = Variable(statement)
        if var.type == VariableType.NORMAL:
            if waiting_variable_msg:
                add_line(waiting_variable_msg, depth)
                waiting_variable_msg = None
                waited_once = False
            add_line(statement, depth)
        elif var.type == VariableType.MSG_REPLY:
            if waiting_variable_msg:
                add_line(waiting_variable_msg, depth) 
            waiting_variable_msg = f'{var.name} = {var.value}'
            waited_once = False
        elif var.type == VariableType.MSG_AUTORET:
            if waiting_variable_msg:
                add_line(var.name + ' = ' + waiting_variable_msg.split(' = ')[1], depth)
                waiting_variable_msg = None
                waited_once = False
            else:
                # processing the previous msg failed
                add_line(statement, depth)
        
        variables[var.name] = var
    else:
        if line not in ['{', '}'] and line:
            if waiting_variable_msg and not waited_once:
                waited_once = True
            if waiting_variable_msg and waited_once:
                add_line(waiting_variable_msg, depth)
                waiting_variable_msg = None
                waited_once = False
            add_line(line.strip(), depth)
        
        
    if line.endswith('{') and len(line.strip()) > 1:
        add_line('{', depth)
        depth += 1
    if line.endswith('}') and len(line.strip()) > 1:
        depth -= 1
        add_line('}', depth)
        
        
finished_lines = []
    
    
cleanup_lines = []


depth = 0

def add_cline(line, depth, indent=4):
    cleanup_lines.append(f'{" " * (indent * depth)}{line}')

indent_next_line = 0  
    
for i, line in enumerate(processed_lines):
    line = line.strip()
    
    if line.startswith('{'):
        add_cline('{', depth)
        depth += 1
    if line.startswith('}'):
        depth -= 1
        add_cline('}', depth)
    if indent_next_line == 1:
        depth+= 1
        indent_next_line += 1
    elif indent_next_line == 2:
        depth-= 1
        indent_next_line = 0
        
    if 'if (' in line and '{' not in processed_lines[i+1]:
        indent_next_line = 1
    elif 'else' in line and '{' not in processed_lines[i+1]:
        indent_next_line = 1
    
    line = line.replace('&OBJC_CLASS___', '')
    if line.strip().startswith('objc_release'):
        continue
    if line.strip().startswith('objc_msgSend('):
        line = proc_objc_msg(line.strip())
    
    if line not in ['{', '}'] and line:
        add_cline(line, depth)
        
for line in cleanup_lines:
    print(line)