Coverage for product_risk_suite/api/schema.py: 79%

1import graphene 

2from graphene_django import DjangoObjectType 

3from guardian.shortcuts import get_objects_for_user 

4 

5from product.models import * 

6from risk_assessment.models import * 

7from threat_model.models import * 

8 

9# PRODUCT 

10class ProductRiskEntryType(DjangoObjectType): 

11 class Meta: 

12 model = ProductRiskEntry 

13 fields = "__all__" 

14 

15class ProductRiskAnalysisType(DjangoObjectType): 

16 class Meta: 

17 model = ProductRiskAnalysis 

18 fields = "__all__" 

19 

20 risk_entries = graphene.List( 

21 ProductRiskEntryType, 

22 svg_id=graphene.Int(required=False), 

23 ) 

24 

25 def resolve_risk_entries(root, info, svg_id=None): 

26 qs = root.risk_entries.all() 

27 if svg_id is not None: 

28 qs = qs.filter(svg_id__id=svg_id) 

29 return qs 

30 

31class ProductType(DjangoObjectType): 

32 class Meta: 

33 model = Product 

34 fields = "__all__" 

35 

36# RISK ASSESSMENT 

37 

38class AssetType(DjangoObjectType): 

39 class Meta: 

40 model = Asset 

41 fields = ("name",) 

42 

43class LifeCycleType(DjangoObjectType): 

44 class Meta: 

45 model = LifeCycle 

46 fields = ("name",) 

47 

48class OriginType(DjangoObjectType): 

49 class Meta: 

50 model = Origin 

51 fields = ("name",) 

52 

53class SecurityRequirementType(DjangoObjectType): 

54 class Meta: 

55 model = SecurityRequirement 

56 exclude = ("id",) 

57 

58class StrideType(DjangoObjectType): 

59 class Meta: 

60 model = Stride 

61 fields = ("name",) 

62 

63 full = graphene.String() 

64 

65 def resolve_full(self, info): 

66 return Stride.STRIDE_CHOICES[self.name] 

67 

68class SuggestedMitigationValidation: 

69 class Meta: 

70 model = SuggestedMitigationValidation 

71 fields = "__all__" 

72 

73class RiskType(DjangoObjectType): 

74 class Meta: 

75 model = Risk 

76 exclude = ("id",) 

77 

78class RiskRatingType(DjangoObjectType): 

79 class Meta: 

80 model = RiskRating 

81 exclude = ("id",) 

82 

83 likelihood_of_occurrence_str = graphene.String() 

84 severity_of_impact_str = graphene.String() 

85 risk = graphene.Int() 

86 

87 def resolve_likelihood_of_occurrence_str(self, info): 

88 return LikelihoodName.to_human_str(self.likelihood_of_occurrence) 

89 

90 def resolve_severity_of_impact_str(self, info): 

91 return SeverityName.to_human_str(self.severity_of_impact) 

92 

93 def resolve_risk(self, info): 

94 return self.risk 

95 

96class StatusType(DjangoObjectType): 

97 class Meta: 

98 model = Status 

99 exclude = ("id",) 

100 

101 is_todo = graphene.Boolean() 

102 

103 def resolve_is_todo(self, info): 

104 return self.is_todo 

105 

106class EvidenceType(DjangoObjectType): 

107 class Meta: 

108 model = Evidence 

109 exclude = ("id",) 

110 

111 # overload to not expose whole user structure 

112 responsible = graphene.String() 

113 

114 def resolve_responsible(self, info): 

115 return str(self.responsible) 

116 

117class RiskMitigationType(DjangoObjectType): 

118 class Meta: 

119 model = RiskMitigation 

120 exclude = ("id",) 

121 

122# Threat Model 

123class ThreatModel(DjangoObjectType): 

124 class Meta: 

125 model = ThreatModel 

126 fields = "__all__" 

127 

128class ThreatModelConnectionNameType(DjangoObjectType): 

129 class Meta: 

130 model = ThreatModelConnectionName 

131 fields = "__all__" 

132 

133# the main query 

134class Query(graphene.ObjectType): 

135 products = graphene.List(ProductType) 

136 product_by_title = graphene.Field(ProductType, title=graphene.String(required=True)) 

137 product_by_slug = graphene.Field(ProductType, slug=graphene.String(required=True)) 

138 

139 security_requirements = graphene.List(SecurityRequirementType) 

140 stride_to_str = graphene.Field(graphene.String, name=graphene.String(required=True)) 

141 

142 def resolve_products(root, info): 

143 if not info.context.user.is_authenticated: 

144 return Product.objects.none() 

145 

146 products = get_objects_for_user(info.context.user, 'product.view_product') 

147 

148 return products.order_by("title") 

149 

150 def resolve_product_by_title(root, info, title): 

151 if not info.context.user.is_authenticated: 

152 return Product.objects.none() 

153 

154 try: 

155 return get_objects_for_user(info.context.user, 'product.view_product').get(title=title) 

156 except Product.DoesNotExist: 

157 return None 

158 

159 def resolve_product_by_slug(root, info, slug): 

160 if not info.context.user.is_authenticated: 

161 return Product.objects.none() 

162 

163 try: 

164 return get_objects_for_user(info.context.user, 'product.view_product').get(slug=slug) 

165 except Product.DoesNotExist: 

166 return None 

167 

168 def resolve_security_requirements(root, info): 

169 return SecurityRequirement.objects.all().order_by("norm_short", "description_short") 

170 

171 def resolve_stride_to_str(root, info, name): 

172 if name in Stride.STRIDE_CHOICES: 

173 return Stride.STRIDE_CHOICES[name] 

174 allowed = [s for s in Stride.STRIDE_CHOICES.keys()] 

175 return f"Invalid name. Allowed: {allowed}" 

176 

177schema = graphene.Schema(query=Query)