如何在不退出程序的情况下检查数字是否为素数并再次询问用户是否为素数？

Question

我试图检查一个数字是否为素数，作为Diffie Hellman算法的一部分，如果它不是素数，则要求用户再次输入另一个数字而不退出程序。

这是我的代码：

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)

Answer 1

user_input = int(input("Enter an integer !  "))

def prime(n):
    isPrime = True
    for num in range(2, int(n/2)):
        if n % num == 0:
            isPrime = False
    return isPrime


while (prime(user_input)== False):
    user_input = int(input("==> Your input is incorrect, please enter  an other integer !"))
print("{} is prime ".format(user_input))