Как проверить, является ли число простым, и снова спросить пользователя, не является ли оно простым, не выходя из программы?

Я пытаюсь проверить, является ли число простым или нет, как часть алгоритма Диффи Хеллмана, и если это не простое число, снова попросите пользователя ввести другое число, не выходя из программы.

Вот мой код:

def prime(p):
    for i in range(2,p):
        if (p % i) == 0:
            return(False)
            break
        else:
            return(True)


    def primRoots(p):
        roots = []
        required_set = set(num for num in range (1, p))

        for g in range(1, p):
            actual_set = set(pow(g, powers) % p for powers in range (1, p))
            if required_set == actual_set:
                roots.append(g)           
        return roots

    p=int(input("enter any prime no:"))
    check=prime(p)
    if(check==True):

        primitive_roots = primRoots(p)

        g=primitive_roots[0]
        print(g)
        x=int(input("Alice chooses value of X as:"))
        y=int(input("Bob chooses value of y as:"))
        r1=(g**x) % p
        r2=(g**y) % p

        print("value of r1 is",r1)
        print("value of r2 is",r2)

        a=x*y

        k1=(r2**x) % p
        print("k1 is",k1)
        k2=(r1**y)% p
        print("k2 is",k2)
        k=(g**a)%p
        print("shared key is",k)    
    else:
        print("It is not a prime num,enter again")
        prime(p)