add 2fa type tracking to credentials (#3647)

This commit is contained in:
pedrohsdb
2025-10-08 11:38:34 -07:00
committed by GitHub
parent a61ff4cb4f
commit fc0f2b87ca
9 changed files with 123 additions and 6 deletions

View File

@@ -0,0 +1,38 @@
"""add_totp_type_to_credentials
Revision ID: 8d015c5358f8
Revises: 81351201bc8d
Create Date: 2025-10-08 18:22:45.893853+00:00
"""
from typing import Sequence, Union
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "8d015c5358f8"
down_revision: Union[str, None] = "81351201bc8d"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
# Add column as nullable initially
op.add_column("credentials", sa.Column("totp_type", sa.String(), nullable=True))
# Set default value for existing records
op.execute("UPDATE credentials SET totp_type = 'none' WHERE totp_type IS NULL")
# Make column non-nullable with default value
op.alter_column("credentials", "totp_type", nullable=False, server_default="none")
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column("credentials", "totp_type")
# ### end Alembic commands ###

View File

@@ -423,6 +423,7 @@ export type Createv2TaskRequest = {
export type PasswordCredentialApiResponse = {
username: string;
totp_type: "authenticator" | "email" | "text" | "none";
};
export type CreditCardCredentialApiResponse = {
@@ -459,6 +460,7 @@ export type PasswordCredential = {
username: string;
password: string;
totp: string | null;
totp_type: "authenticator" | "email" | "text" | "none";
};
export type CreditCardCredential = {

View File

@@ -7,6 +7,20 @@ type Props = {
};
function CredentialItem({ credential }: Props) {
const getTotpTypeDisplay = (totpType: string) => {
switch (totpType) {
case "authenticator":
return "Authenticator App";
case "email":
return "Email";
case "text":
return "Text Message";
case "none":
default:
return "";
}
};
return (
<div className="flex gap-5 rounded-lg bg-slate-elevation2 p-4">
<div className="w-48 space-y-2">
@@ -21,10 +35,18 @@ function CredentialItem({ credential }: Props) {
<div className="shrink-0 space-y-2">
<p className="text-sm text-slate-400">Username/Email</p>
<p className="text-sm text-slate-400">Password</p>
{credential.credential.totp_type !== "none" && (
<p className="text-sm text-slate-400">2FA Type</p>
)}
</div>
<div className="space-y-2">
<p className="text-sm">{credential.credential.username}</p>
<p className="text-sm">{"********"}</p>
{credential.credential.totp_type !== "none" && (
<p className="text-sm text-blue-400">
{getTotpTypeDisplay(credential.credential.totp_type)}
</p>
)}
</div>
</div>
</div>

View File

@@ -27,6 +27,7 @@ const PASSWORD_CREDENTIAL_INITIAL_VALUES = {
username: "",
password: "",
totp: "",
totp_type: "none" as "none" | "authenticator" | "email" | "text",
};
const CREDIT_CARD_CREDENTIAL_INITIAL_VALUES = {
@@ -156,6 +157,7 @@ function CredentialsModal({ onCredentialCreated }: Props) {
username,
password,
totp: totp === "" ? null : totp,
totp_type: passwordCredentialValues.totp_type,
},
});
} else if (type === CredentialModalTypes.CREDIT_CARD) {

View File

@@ -24,24 +24,40 @@ type Props = {
username: string;
password: string;
totp: string;
totp_type: "authenticator" | "email" | "text" | "none";
};
onChange: (values: {
name: string;
username: string;
password: string;
totp: string;
totp_type: "authenticator" | "email" | "text" | "none";
}) => void;
};
function PasswordCredentialContent({
values: { name, username, password, totp },
values: { name, username, password, totp, totp_type },
onChange,
}: Props) {
const [totpMethod, setTotpMethod] = useState<
"text" | "email" | "authenticator"
"authenticator" | "email" | "text"
>("authenticator");
const [showPassword, setShowPassword] = useState(false);
// Update totp_type when totpMethod changes
const handleTotpMethodChange = (
method: "authenticator" | "email" | "text",
) => {
setTotpMethod(method);
onChange({
name,
username,
password,
totp: method === "authenticator" ? totp : "",
totp_type: method,
});
};
return (
<div className="space-y-5">
<div className="flex">
@@ -59,6 +75,7 @@ function PasswordCredentialContent({
username,
password,
totp,
totp_type,
})
}
/>
@@ -76,6 +93,7 @@ function PasswordCredentialContent({
username: e.target.value,
password,
totp,
totp_type,
})
}
/>
@@ -95,6 +113,7 @@ function PasswordCredentialContent({
username,
totp,
password: e.target.value,
totp_type,
})
}
/>
@@ -133,7 +152,7 @@ function PasswordCredentialContent({
"bg-slate-elevation3": totpMethod === "authenticator",
},
)}
onClick={() => setTotpMethod("authenticator")}
onClick={() => handleTotpMethodChange("authenticator")}
>
<QRCodeIcon className="h-6 w-6" />
<Label>Authenticator App</Label>
@@ -145,7 +164,7 @@ function PasswordCredentialContent({
"bg-slate-elevation3": totpMethod === "email",
},
)}
onClick={() => setTotpMethod("email")}
onClick={() => handleTotpMethodChange("email")}
>
<EnvelopeClosedIcon className="h-6 w-6" />
<Label>Email</Label>
@@ -157,7 +176,7 @@ function PasswordCredentialContent({
"bg-slate-elevation3": totpMethod === "text",
},
)}
onClick={() => setTotpMethod("text")}
onClick={() => handleTotpMethodChange("text")}
>
<MobileIcon className="h-6 w-6" />
<Label>Text Message</Label>
@@ -202,6 +221,7 @@ function PasswordCredentialContent({
username,
password,
totp: e.target.value,
totp_type,
})
}
/>

View File

@@ -3623,6 +3623,7 @@ class AgentDB:
credential_type: CredentialType,
organization_id: str,
item_id: str,
totp_type: str = "none",
) -> Credential:
async with self.Session() as session:
credential = CredentialModel(
@@ -3630,6 +3631,7 @@ class AgentDB:
name=name,
credential_type=credential_type,
item_id=item_id,
totp_type=totp_type,
)
session.add(credential)
await session.commit()

View File

@@ -818,6 +818,7 @@ class CredentialModel(Base):
name = Column(String, nullable=False)
credential_type = Column(String, nullable=False)
totp_type = Column(String, nullable=False, default="none")
created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)

View File

@@ -167,11 +167,13 @@ async def create_credential(
item_id=item_id,
name=data.name,
credential_type=data.credential_type,
totp_type=data.credential.totp_type if hasattr(data.credential, "totp_type") else "none",
)
if data.credential_type == CredentialType.PASSWORD:
credential_response = PasswordCredentialResponse(
username=data.credential.username,
totp_type=data.credential.totp_type if hasattr(data.credential, "totp_type") else "none",
)
return CredentialResponse(
credential=credential_response,
@@ -283,6 +285,7 @@ async def get_credential(
if credential_item.credential_type == CredentialType.PASSWORD:
credential_response = PasswordCredentialResponse(
username=credential_item.credential.username,
totp_type=credential.totp_type,
)
return CredentialResponse(
credential=credential_response,
@@ -354,7 +357,10 @@ async def get_credentials(
if not item:
continue
if item.credential_type == CredentialType.PASSWORD:
credential_response = PasswordCredentialResponse(username=item.credential.username)
credential_response = PasswordCredentialResponse(
username=item.credential.username,
totp_type=credential.totp_type,
)
response_items.append(
CredentialResponse(
credential=credential_response,

View File

@@ -11,10 +11,24 @@ class CredentialType(StrEnum):
CREDIT_CARD = "credit_card"
class TotpType(StrEnum):
"""Type of 2FA/TOTP method used."""
AUTHENTICATOR = "authenticator"
EMAIL = "email"
TEXT = "text"
NONE = "none"
class PasswordCredentialResponse(BaseModel):
"""Response model for password credentials, containing only the username."""
username: str = Field(..., description="The username associated with the credential", examples=["user@example.com"])
totp_type: TotpType = Field(
TotpType.NONE,
description="Type of 2FA method used for this credential",
examples=[TotpType.AUTHENTICATOR],
)
class CreditCardCredentialResponse(BaseModel):
@@ -34,6 +48,11 @@ class PasswordCredential(BaseModel):
description="Optional TOTP (Time-based One-Time Password) string used to generate 2FA codes",
examples=["JBSWY3DPEHPK3PXP"],
)
totp_type: TotpType = Field(
TotpType.NONE,
description="Type of 2FA method used for this credential",
examples=[TotpType.AUTHENTICATOR],
)
class NonEmptyPasswordCredential(PasswordCredential):
@@ -124,6 +143,11 @@ class Credential(BaseModel):
name: str = Field(..., description="Name of the credential", examples=["Skyvern Login"])
credential_type: CredentialType = Field(..., description="Type of the credential. Eg password, credit card, etc.")
item_id: str = Field(..., description="ID of the associated credential item", examples=["item_1234567890"])
totp_type: TotpType = Field(
TotpType.NONE,
description="Type of 2FA method used for this credential",
examples=[TotpType.AUTHENTICATOR],
)
created_at: datetime = Field(..., description="Timestamp when the credential was created")
modified_at: datetime = Field(..., description="Timestamp when the credential was last modified")