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

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))